<commit_before>use cocoa::base::id;

/// The `MTLLibrary` protocol defines the interface for an object that represents a library of Metal
/// shader functions. A `MTLLibrary` object can contain Metal shading language code that is compiled
/// during the app build process or at runtime from a text string containing Metal shading language
/// source code. Your app does not define classes that implement this protocol.
///
/// Use a `MTLDevice` method (not standard allocation and initialization techniques) to create a
/// `MTLLibrary` object. To create a `MTLLibrary` object from a Metal library binary, call one of
/// these MTLDevice methods:
///
/// * `newDefaultLibrary`
///
/// * `newLibraryWithFile:error:`
///
/// * `newLibraryWithData:error:`
///
/// To create a `MTLLibrary` object by compiling source code, call one of these `MTLDevice` methods:
///
/// * `newLibraryWithSource:options:error:`
///
/// * `newLibraryWithSource:options:completionHandler:`
///
/// The `newFunctionWithName:` method is used to fetch functions from the library, which makes that
/// code available as a shader for either a `MTLRenderPipelineState` object for a render command
/// encoder or for a `MTLComputePipelineState` for a compute command encoder.
pub trait MTLLibrary {
    /// Returns a function object that represents an entry point in the library.
    ///
    /// # Parameters
    ///
    /// * `functionName` - The name of an entry point.
    ///
    /// # Return Value
    ///
    /// A function object for the named entry point, or nil if the named function is not found in
    /// the library.
    unsafe fn newFunctionWithName(self, functionName: id) -> id;

    /// A list of all entry points in the library. (read-only)
    ///
    /// # Discussion
    ///
    /// An array of `NSString` objects. Each string is the name of an entry point.
    unsafe fn functionNames(self) -> id;

    /// The device from which this library was created. (read-only)
    ///
    /// # Discussion
    ///
    /// This library can only be used with this device.
    unsafe fn device(self) -> id;

    /// A string to help identify the library object.
    unsafe fn label(self) -> id;
    unsafe fn setLabel(self, id);
}

impl MTLLibrary for id {
    unsafe fn newFunctionWithName(self, functionName: id) -> id {
        msg_send![self, newFunctionWithName:functionName]
    }

    unsafe fn functionNames(self) -> id {
        msg_send![self, functionNames]
    }

    unsafe fn device(self) -> id {
        msg_send![self, device]
    }

    unsafe fn label(self) -> id {
        msg_send![self, label]
    }

    unsafe fn setLabel(self, label: id) {
        msg_send![self, setLabel:label]
    }
}

/// Error conditions that can result from the creation of a `MTLLibrary` or `MTLFunction` object.
#[repr(C, usize)]
#[derive(Clone, Copy, Eq, Hash, PartialEq)]
pub enum MTLLibraryError {
    /// The action is unsupported. For example, the requested library file has improper formatting,
    /// or the requested library is not accessible.
    MTLLibraryErrorUnsupported = 1,

    /// The action causes an internal error.
    MTLLibraryErrorInternal = 2,

    /// Compilation fails
    MTLLibraryErrorCompileFailure = 3,

    /// Compilation succeeds without error, but there are compiler warnings.
    MTLLibraryErrorCompileWarning = 4
}

#[link(name = "Metal", kind = "framework")]
extern "C" {
    /// Constant to identify the `MTLLibrary` error domain.
    pub static MTLLibraryErrorDomain: id;
}

/// Error conditions that can result from the creation of a `MTLRenderPipelineState` object.
#[repr(C, usize)]
#[derive(Clone, Copy, Eq, Hash, PartialEq)]
pub enum MTLRenderPipelineError {
    /// The action causes an internal error.
    MTLRenderPipelineErrorInternal = 1,

    /// The action is unsupported.
    MTLRenderPipelineErrorUnsupported = 2,

    /// The input values are invalid.
    MTLRenderPipelineErrorInvalidInput = 3
}

#[link(name = "Metal", kind = "framework")]
extern "C" {
    /// Constant to identify the `MTLRenderPipelineState` error domain.
    pub static MTLRenderPipelineErrorDomain: id;
}

#[repr(C, usize)]
#[derive(Clone, Copy, Eq, Hash, PartialEq)]
pub enum MTLLanguageVersion {
    /// Version 1.0
    MTLLanguageVersion1_0 = (1 << 16),

    /// Version 1.1
    MTLLanguageVersion1_1 = (1 << 16) + 1
}
<commit_msg>Add derive(Debug) for some enums<commit_after>use cocoa::base::id;

/// The `MTLLibrary` protocol defines the interface for an object that represents a library of Metal
/// shader functions. A `MTLLibrary` object can contain Metal shading language code that is compiled
/// during the app build process or at runtime from a text string containing Metal shading language
/// source code. Your app does not define classes that implement this protocol.
///
/// Use a `MTLDevice` method (not standard allocation and initialization techniques) to create a
/// `MTLLibrary` object. To create a `MTLLibrary` object from a Metal library binary, call one of
/// these MTLDevice methods:
///
/// * `newDefaultLibrary`
///
/// * `newLibraryWithFile:error:`
///
/// * `newLibraryWithData:error:`
///
/// To create a `MTLLibrary` object by compiling source code, call one of these `MTLDevice` methods:
///
/// * `newLibraryWithSource:options:error:`
///
/// * `newLibraryWithSource:options:completionHandler:`
///
/// The `newFunctionWithName:` method is used to fetch functions from the library, which makes that
/// code available as a shader for either a `MTLRenderPipelineState` object for a render command
/// encoder or for a `MTLComputePipelineState` for a compute command encoder.
pub trait MTLLibrary {
    /// Returns a function object that represents an entry point in the library.
    ///
    /// # Parameters
    ///
    /// * `functionName` - The name of an entry point.
    ///
    /// # Return Value
    ///
    /// A function object for the named entry point, or nil if the named function is not found in
    /// the library.
    unsafe fn newFunctionWithName(self, functionName: id) -> id;

    /// A list of all entry points in the library. (read-only)
    ///
    /// # Discussion
    ///
    /// An array of `NSString` objects. Each string is the name of an entry point.
    unsafe fn functionNames(self) -> id;

    /// The device from which this library was created. (read-only)
    ///
    /// # Discussion
    ///
    /// This library can only be used with this device.
    unsafe fn device(self) -> id;

    /// A string to help identify the library object.
    unsafe fn label(self) -> id;
    unsafe fn setLabel(self, id);
}

impl MTLLibrary for id {
    unsafe fn newFunctionWithName(self, functionName: id) -> id {
        msg_send![self, newFunctionWithName:functionName]
    }

    unsafe fn functionNames(self) -> id {
        msg_send![self, functionNames]
    }

    unsafe fn device(self) -> id {
        msg_send![self, device]
    }

    unsafe fn label(self) -> id {
        msg_send![self, label]
    }

    unsafe fn setLabel(self, label: id) {
        msg_send![self, setLabel:label]
    }
}

/// Error conditions that can result from the creation of a `MTLLibrary` or `MTLFunction` object.
#[repr(C, usize)]
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum MTLLibraryError {
    /// The action is unsupported. For example, the requested library file has improper formatting,
    /// or the requested library is not accessible.
    MTLLibraryErrorUnsupported = 1,

    /// The action causes an internal error.
    MTLLibraryErrorInternal = 2,

    /// Compilation fails
    MTLLibraryErrorCompileFailure = 3,

    /// Compilation succeeds without error, but there are compiler warnings.
    MTLLibraryErrorCompileWarning = 4
}

#[link(name = "Metal", kind = "framework")]
extern "C" {
    /// Constant to identify the `MTLLibrary` error domain.
    pub static MTLLibraryErrorDomain: id;
}

/// Error conditions that can result from the creation of a `MTLRenderPipelineState` object.
#[repr(C, usize)]
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum MTLRenderPipelineError {
    /// The action causes an internal error.
    MTLRenderPipelineErrorInternal = 1,

    /// The action is unsupported.
    MTLRenderPipelineErrorUnsupported = 2,

    /// The input values are invalid.
    MTLRenderPipelineErrorInvalidInput = 3
}

#[link(name = "Metal", kind = "framework")]
extern "C" {
    /// Constant to identify the `MTLRenderPipelineState` error domain.
    pub static MTLRenderPipelineErrorDomain: id;
}

#[repr(C, usize)]
#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
pub enum MTLLanguageVersion {
    /// Version 1.0
    MTLLanguageVersion1_0 = (1 << 16),

    /// Version 1.1
    MTLLanguageVersion1_1 = (1 << 16) + 1
}
<|endoftext|>


<commit_before>#![cfg(test)]
use super::*;
use crate::test_utils::*;
use std::ops::RangeInclusive;

ssz_tests!(FoundationBeaconState);

fn test_beacon_proposer_index<T: EthSpec>() {
    let spec = T::default_spec();
    let relative_epoch = RelativeEpoch::Current;

    // Build a state for testing.
    let build_state = |validator_count: usize| -> BeaconState<T> {
        let builder: TestingBeaconStateBuilder<T> =
            TestingBeaconStateBuilder::from_default_keypairs_file_if_exists(validator_count, &spec);
        let (mut state, _keypairs) = builder.build();
        state.build_committee_cache(relative_epoch, &spec).unwrap();

        state
    };

    // Run a test on the state.
    let test = |state: &BeaconState<T>, slot: Slot, shuffling_index: usize| {
        let shuffling = state.get_shuffling(relative_epoch).unwrap();
        assert_eq!(
            state.get_beacon_proposer_index(slot, relative_epoch, &spec),
            Ok(shuffling[shuffling_index])
        );
    };

    // Test where we have one validator per slot
    let state = build_state(T::slots_per_epoch() as usize);
    for i in 0..T::slots_per_epoch() {
        test(&state, Slot::from(i), i as usize);
    }

    // Test where we have two validators per slot
    let state = build_state(T::slots_per_epoch() as usize * 2);
    for i in 0..T::slots_per_epoch() {
        test(&state, Slot::from(i), i as usize * 2);
    }

    // Test with two validators per slot, first validator has zero balance.
    let mut state = build_state(T::slots_per_epoch() as usize * 2);
    let shuffling = state.get_shuffling(relative_epoch).unwrap().to_vec();
    state.validators[shuffling[0]].effective_balance = 0;
    test(&state, Slot::new(0), 1);
    for i in 1..T::slots_per_epoch() {
        test(&state, Slot::from(i), i as usize * 2);
    }
}

#[test]
fn beacon_proposer_index() {
    test_beacon_proposer_index::<MinimalEthSpec>();
}

/// Should produce (note the set notation brackets):
///
/// (current_epoch - LATEST_ACTIVE_INDEX_ROOTS_LENGTH + ACTIVATION_EXIT_DELAY, current_epoch +
/// ACTIVATION_EXIT_DELAY]
fn active_index_range<T: EthSpec>(current_epoch: Epoch) -> RangeInclusive<Epoch> {
    let delay = T::default_spec().activation_exit_delay;

    let start: i32 =
        current_epoch.as_u64() as i32 - T::epochs_per_historical_vector() as i32 + delay as i32;
    let end = current_epoch + delay;

    let start: Epoch = if start < 0 {
        Epoch::new(0)
    } else {
        Epoch::from(start as u64 + 1)
    };

    start..=end
}

/// Test getting an active index root at the start and end of the valid range, and one either side
/// of that range.
fn test_active_index<T: EthSpec>(state_slot: Slot) {
    let spec = T::default_spec();
    let builder: TestingBeaconStateBuilder<T> =
        TestingBeaconStateBuilder::from_default_keypairs_file_if_exists(16, &spec);
    let (mut state, _keypairs) = builder.build();
    state.slot = state_slot;

    let range = active_index_range::<T>(state.current_epoch());

    let modulo = |epoch: Epoch| epoch.as_usize() % T::epochs_per_historical_vector();

    // Test the start and end of the range.
    assert_eq!(
        state.get_active_index_root_index(*range.start(), &spec),
        Ok(modulo(*range.start()))
    );
    assert_eq!(
        state.get_active_index_root_index(*range.end(), &spec),
        Ok(modulo(*range.end()))
    );

    // One either side of the range.
    if state.current_epoch() > 0 {
        // Test is invalid on epoch zero, cannot subtract from zero.
        assert_eq!(
            state.get_active_index_root_index(*range.start() - 1, &spec),
            Err(Error::EpochOutOfBounds)
        );
    }
    assert_eq!(
        state.get_active_index_root_index(*range.end() + 1, &spec),
        Err(Error::EpochOutOfBounds)
    );
}

#[test]
fn get_active_index_root_index() {
    test_active_index::<MainnetEthSpec>(Slot::new(0));

    let epoch = Epoch::from(MainnetEthSpec::epochs_per_historical_vector() * 4);
    let slot = epoch.start_slot(MainnetEthSpec::slots_per_epoch());
    test_active_index::<MainnetEthSpec>(slot);
}

/// Test that
///
/// 1. Using the cache before it's built fails.
/// 2. Using the cache after it's build passes.
/// 3. Using the cache after it's dropped fails.
fn test_cache_initialization<'a, T: EthSpec>(
    state: &'a mut BeaconState<T>,
    relative_epoch: RelativeEpoch,
    spec: &ChainSpec,
) {
    let slot = relative_epoch
        .into_epoch(state.slot.epoch(T::slots_per_epoch()))
        .start_slot(T::slots_per_epoch());

    // Assuming the cache isn't already built, assert that a call to a cache-using function fails.
    assert_eq!(
        state.get_attestation_duties(0, relative_epoch),
        Err(BeaconStateError::CommitteeCacheUninitialized(
            relative_epoch
        ))
    );

    // Build the cache.
    state.build_committee_cache(relative_epoch, spec).unwrap();

    // Assert a call to a cache-using function passes.
    let _ = state
        .get_beacon_proposer_index(slot, relative_epoch, spec)
        .unwrap();

    // Drop the cache.
    state.drop_committee_cache(relative_epoch);

    // Assert a call to a cache-using function fail.
    assert_eq!(
        state.get_beacon_proposer_index(slot, relative_epoch, spec),
        Err(BeaconStateError::CommitteeCacheUninitialized(
            relative_epoch
        ))
    );
}

#[test]
fn cache_initialization() {
    let spec = MinimalEthSpec::default_spec();

    let builder: TestingBeaconStateBuilder<MinimalEthSpec> =
        TestingBeaconStateBuilder::from_default_keypairs_file_if_exists(16, &spec);
    let (mut state, _keypairs) = builder.build();

    state.slot =
        (MinimalEthSpec::genesis_epoch() + 1).start_slot(MinimalEthSpec::slots_per_epoch());

    test_cache_initialization(&mut state, RelativeEpoch::Previous, &spec);
    test_cache_initialization(&mut state, RelativeEpoch::Current, &spec);
    test_cache_initialization(&mut state, RelativeEpoch::Next, &spec);
}

#[test]
fn tree_hash_cache() {
    use crate::test_utils::{SeedableRng, TestRandom, XorShiftRng};
    use tree_hash::TreeHash;

    let mut rng = XorShiftRng::from_seed([42; 16]);

    let mut state: FoundationBeaconState = BeaconState::random_for_test(&mut rng);

    let root = state.update_tree_hash_cache().unwrap();

    assert_eq!(root.as_bytes(), &state.tree_hash_root()[..]);

    state.slot = state.slot + 1;

    let root = state.update_tree_hash_cache().unwrap();
    assert_eq!(root.as_bytes(), &state.tree_hash_root()[..]);
}

/// Tests committee-specific components
#[cfg(test)]
mod committees {
    use super::*;
    use crate::beacon_state::MinimalEthSpec;
    use swap_or_not_shuffle::shuffle_list;

    fn execute_committee_consistency_test<T: EthSpec>(
        state: BeaconState<T>,
        epoch: Epoch,
        validator_count: usize,
        spec: &ChainSpec,
    ) {
        let active_indices: Vec<usize> = (0..validator_count).collect();
        let seed = state.get_seed(epoch, spec).unwrap();
        let relative_epoch = RelativeEpoch::from_epoch(state.current_epoch(), epoch).unwrap();
        let start_shard =
            CommitteeCache::compute_start_shard(&state, relative_epoch, active_indices.len(), spec);

        let mut ordered_indices = state
            .get_cached_active_validator_indices(relative_epoch)
            .unwrap()
            .to_vec();
        ordered_indices.sort_unstable();
        assert_eq!(
            active_indices, ordered_indices,
            "Validator indices mismatch"
        );

        let shuffling =
            shuffle_list(active_indices, spec.shuffle_round_count, &seed[..], false).unwrap();

        let mut expected_indices_iter = shuffling.iter();
        let mut expected_shards_iter = (0..T::ShardCount::to_u64())
            .into_iter()
            .map(|i| (start_shard + i) % T::ShardCount::to_u64());

        // Loop through all slots in the epoch being tested.
        for slot in epoch.slot_iter(T::slots_per_epoch()) {
            let crosslink_committees = state.get_crosslink_committees_at_slot(slot).unwrap();

            // Assert that the number of committees in this slot is consistent with the reported number
            // of committees in an epoch.
            assert_eq!(
                crosslink_committees.len() as u64,
                state.get_committee_count(relative_epoch).unwrap() / T::slots_per_epoch()
            );

            for cc in crosslink_committees {
                // Assert that shards are assigned contiguously across committees.
                assert_eq!(expected_shards_iter.next().unwrap(), cc.shard);
                // Assert that a committee lookup via slot is identical to a committee lookup via
                // shard.
                assert_eq!(
                    state
                        .get_crosslink_committee_for_shard(cc.shard, relative_epoch)
                        .unwrap(),
                    cc
                );

                // Loop through each validator in the committee.
                for (committee_i, validator_i) in cc.committee.iter().enumerate() {
                    // Assert the validators are assigned contiguously across committees.
                    assert_eq!(
                        *validator_i,
                        *expected_indices_iter.next().unwrap(),
                        "Non-sequential validators."
                    );
                    // Assert a call to `get_attestation_duties` is consistent with a call to
                    // `get_crosslink_committees_at_slot`
                    let attestation_duty = state
                        .get_attestation_duties(*validator_i, relative_epoch)
                        .unwrap()
                        .unwrap();
                    assert_eq!(attestation_duty.slot, slot);
                    assert_eq!(attestation_duty.shard, cc.shard);
                    assert_eq!(attestation_duty.committee_index, committee_i);
                    assert_eq!(attestation_duty.committee_len, cc.committee.len());
                }
            }
        }

        // Assert that all validators were assigned to a committee.
        assert!(expected_indices_iter.next().is_none());

        // Assert that all shards were assigned to a committee.
        assert!(expected_shards_iter.next().is_none());
    }

    fn committee_consistency_test<T: EthSpec>(
        validator_count: usize,
        state_epoch: Epoch,
        cache_epoch: RelativeEpoch,
    ) {
        let spec = &T::default_spec();

        let mut builder = TestingBeaconStateBuilder::from_single_keypair(
            validator_count,
            &Keypair::random(),
            spec,
        );

        let slot = state_epoch.start_slot(T::slots_per_epoch());
        builder.teleport_to_slot(slot);

        let (mut state, _keypairs): (BeaconState<T>, _) = builder.build();

        let distinct_hashes: Vec<Hash256> = (0..T::epochs_per_historical_vector())
            .into_iter()
            .map(|i| Hash256::from_low_u64_be(i as u64))
            .collect();
        state.randao_mixes = FixedVector::from(distinct_hashes);

        state
            .build_committee_cache(RelativeEpoch::Previous, spec)
            .unwrap();
        state
            .build_committee_cache(RelativeEpoch::Current, spec)
            .unwrap();
        state
            .build_committee_cache(RelativeEpoch::Next, spec)
            .unwrap();

        let cache_epoch = cache_epoch.into_epoch(state_epoch);

        execute_committee_consistency_test(state, cache_epoch, validator_count as usize, &spec);
    }

    fn committee_consistency_test_suite<T: EthSpec>(cached_epoch: RelativeEpoch) {
        let spec = T::default_spec();

        let validator_count = (T::shard_count() * spec.target_committee_size) + 1;

        committee_consistency_test::<T>(validator_count as usize, Epoch::new(0), cached_epoch);

        committee_consistency_test::<T>(
            validator_count as usize,
            T::genesis_epoch() + 4,
            cached_epoch,
        );

        committee_consistency_test::<T>(
            validator_count as usize,
            T::genesis_epoch() + T::slots_per_historical_root() as u64 * T::slots_per_epoch() * 4,
            cached_epoch,
        );
    }

    #[test]
    fn current_epoch_committee_consistency() {
        committee_consistency_test_suite::<MinimalEthSpec>(RelativeEpoch::Current);
    }

    #[test]
    fn previous_epoch_committee_consistency() {
        committee_consistency_test_suite::<MinimalEthSpec>(RelativeEpoch::Previous);
    }

    #[test]
    fn next_epoch_committee_consistency() {
        committee_consistency_test_suite::<MinimalEthSpec>(RelativeEpoch::Next);
    }
}
<commit_msg>Fix BeaconChain tests<commit_after>#![cfg(test)]
use super::*;
use crate::test_utils::*;
use std::ops::RangeInclusive;

ssz_tests!(FoundationBeaconState);

fn test_beacon_proposer_index<T: EthSpec>() {
    let spec = T::default_spec();
    let relative_epoch = RelativeEpoch::Current;

    // Build a state for testing.
    let build_state = |validator_count: usize| -> BeaconState<T> {
        let builder: TestingBeaconStateBuilder<T> =
            TestingBeaconStateBuilder::from_default_keypairs_file_if_exists(validator_count, &spec);
        let (mut state, _keypairs) = builder.build();
        state.build_committee_cache(relative_epoch, &spec).unwrap();

        state
    };

    // Run a test on the state.
    let test = |state: &BeaconState<T>, slot: Slot, shuffling_index: usize| {
        let shuffling = state.get_shuffling(relative_epoch).unwrap();
        assert_eq!(
            state.get_beacon_proposer_index(slot, relative_epoch, &spec),
            Ok(shuffling[shuffling_index])
        );
    };

    // Test where we have one validator per slot
    let state = build_state(T::slots_per_epoch() as usize);
    for i in 0..T::slots_per_epoch() {
        test(&state, Slot::from(i), i as usize);
    }

    // Test where we have two validators per slot
    let state = build_state(T::slots_per_epoch() as usize * 2);
    for i in 0..T::slots_per_epoch() {
        test(&state, Slot::from(i), i as usize * 2);
    }

    // Test with two validators per slot, first validator has zero balance.
    let mut state = build_state(T::slots_per_epoch() as usize * 2);
    let shuffling = state.get_shuffling(relative_epoch).unwrap().to_vec();
    state.validators[shuffling[0]].effective_balance = 0;
    test(&state, Slot::new(0), 1);
    for i in 1..T::slots_per_epoch() {
        test(&state, Slot::from(i), i as usize * 2);
    }
}

#[test]
fn beacon_proposer_index() {
    test_beacon_proposer_index::<MinimalEthSpec>();
}

/// Should produce (note the set notation brackets):
///
/// (current_epoch - LATEST_ACTIVE_INDEX_ROOTS_LENGTH + ACTIVATION_EXIT_DELAY, current_epoch +
/// ACTIVATION_EXIT_DELAY]
fn active_index_range<T: EthSpec>(current_epoch: Epoch) -> RangeInclusive<Epoch> {
    let delay = T::default_spec().activation_exit_delay;

    let start: i32 =
        current_epoch.as_u64() as i32 - T::epochs_per_historical_vector() as i32 + delay as i32;
    let end = current_epoch + delay;

    let start: Epoch = if start < 0 {
        Epoch::new(0)
    } else {
        Epoch::from(start as u64 + 1)
    };

    start..=end
}

/// Test getting an active index root at the start and end of the valid range, and one either side
/// of that range.
fn test_active_index<T: EthSpec>(state_slot: Slot) {
    let spec = T::default_spec();
    let builder: TestingBeaconStateBuilder<T> =
        TestingBeaconStateBuilder::from_default_keypairs_file_if_exists(16, &spec);
    let (mut state, _keypairs) = builder.build();
    state.slot = state_slot;

    let range = active_index_range::<T>(state.current_epoch());

    let modulo = |epoch: Epoch| epoch.as_usize() % T::epochs_per_historical_vector();

    // Test the start and end of the range.
    assert_eq!(
        state.get_active_index_root_index(*range.start(), &spec, AllowNextEpoch::False),
        Ok(modulo(*range.start()))
    );
    assert_eq!(
        state.get_active_index_root_index(*range.end(), &spec, AllowNextEpoch::False),
        Ok(modulo(*range.end()))
    );

    // One either side of the range.
    if state.current_epoch() > 0 {
        // Test is invalid on epoch zero, cannot subtract from zero.
        assert_eq!(
            state.get_active_index_root_index(*range.start() - 1, &spec, AllowNextEpoch::False),
            Err(Error::EpochOutOfBounds)
        );
    }
    assert_eq!(
        state.get_active_index_root_index(*range.end() + 1, &spec, AllowNextEpoch::False),
        Err(Error::EpochOutOfBounds)
    );
}

#[test]
fn get_active_index_root_index() {
    test_active_index::<MainnetEthSpec>(Slot::new(0));

    let epoch = Epoch::from(MainnetEthSpec::epochs_per_historical_vector() * 4);
    let slot = epoch.start_slot(MainnetEthSpec::slots_per_epoch());
    test_active_index::<MainnetEthSpec>(slot);
}

/// Test that
///
/// 1. Using the cache before it's built fails.
/// 2. Using the cache after it's build passes.
/// 3. Using the cache after it's dropped fails.
fn test_cache_initialization<'a, T: EthSpec>(
    state: &'a mut BeaconState<T>,
    relative_epoch: RelativeEpoch,
    spec: &ChainSpec,
) {
    let slot = relative_epoch
        .into_epoch(state.slot.epoch(T::slots_per_epoch()))
        .start_slot(T::slots_per_epoch());

    // Assuming the cache isn't already built, assert that a call to a cache-using function fails.
    assert_eq!(
        state.get_attestation_duties(0, relative_epoch),
        Err(BeaconStateError::CommitteeCacheUninitialized(
            relative_epoch
        ))
    );

    // Build the cache.
    state.build_committee_cache(relative_epoch, spec).unwrap();

    // Assert a call to a cache-using function passes.
    let _ = state
        .get_beacon_proposer_index(slot, relative_epoch, spec)
        .unwrap();

    // Drop the cache.
    state.drop_committee_cache(relative_epoch);

    // Assert a call to a cache-using function fail.
    assert_eq!(
        state.get_beacon_proposer_index(slot, relative_epoch, spec),
        Err(BeaconStateError::CommitteeCacheUninitialized(
            relative_epoch
        ))
    );
}

#[test]
fn cache_initialization() {
    let spec = MinimalEthSpec::default_spec();

    let builder: TestingBeaconStateBuilder<MinimalEthSpec> =
        TestingBeaconStateBuilder::from_default_keypairs_file_if_exists(16, &spec);
    let (mut state, _keypairs) = builder.build();

    state.slot =
        (MinimalEthSpec::genesis_epoch() + 1).start_slot(MinimalEthSpec::slots_per_epoch());

    test_cache_initialization(&mut state, RelativeEpoch::Previous, &spec);
    test_cache_initialization(&mut state, RelativeEpoch::Current, &spec);
    test_cache_initialization(&mut state, RelativeEpoch::Next, &spec);
}

#[test]
fn tree_hash_cache() {
    use crate::test_utils::{SeedableRng, TestRandom, XorShiftRng};
    use tree_hash::TreeHash;

    let mut rng = XorShiftRng::from_seed([42; 16]);

    let mut state: FoundationBeaconState = BeaconState::random_for_test(&mut rng);

    let root = state.update_tree_hash_cache().unwrap();

    assert_eq!(root.as_bytes(), &state.tree_hash_root()[..]);

    state.slot = state.slot + 1;

    let root = state.update_tree_hash_cache().unwrap();
    assert_eq!(root.as_bytes(), &state.tree_hash_root()[..]);
}

/// Tests committee-specific components
#[cfg(test)]
mod committees {
    use super::*;
    use crate::beacon_state::MinimalEthSpec;
    use swap_or_not_shuffle::shuffle_list;

    fn execute_committee_consistency_test<T: EthSpec>(
        state: BeaconState<T>,
        epoch: Epoch,
        validator_count: usize,
        spec: &ChainSpec,
    ) {
        let active_indices: Vec<usize> = (0..validator_count).collect();
        let seed = state.get_seed(epoch, spec).unwrap();
        let relative_epoch = RelativeEpoch::from_epoch(state.current_epoch(), epoch).unwrap();
        let start_shard =
            CommitteeCache::compute_start_shard(&state, relative_epoch, active_indices.len(), spec);

        let mut ordered_indices = state
            .get_cached_active_validator_indices(relative_epoch)
            .unwrap()
            .to_vec();
        ordered_indices.sort_unstable();
        assert_eq!(
            active_indices, ordered_indices,
            "Validator indices mismatch"
        );

        let shuffling =
            shuffle_list(active_indices, spec.shuffle_round_count, &seed[..], false).unwrap();

        let mut expected_indices_iter = shuffling.iter();
        let mut expected_shards_iter = (0..T::ShardCount::to_u64())
            .into_iter()
            .map(|i| (start_shard + i) % T::ShardCount::to_u64());

        // Loop through all slots in the epoch being tested.
        for slot in epoch.slot_iter(T::slots_per_epoch()) {
            let crosslink_committees = state.get_crosslink_committees_at_slot(slot).unwrap();

            // Assert that the number of committees in this slot is consistent with the reported number
            // of committees in an epoch.
            assert_eq!(
                crosslink_committees.len() as u64,
                state.get_committee_count(relative_epoch).unwrap() / T::slots_per_epoch()
            );

            for cc in crosslink_committees {
                // Assert that shards are assigned contiguously across committees.
                assert_eq!(expected_shards_iter.next().unwrap(), cc.shard);
                // Assert that a committee lookup via slot is identical to a committee lookup via
                // shard.
                assert_eq!(
                    state
                        .get_crosslink_committee_for_shard(cc.shard, relative_epoch)
                        .unwrap(),
                    cc
                );

                // Loop through each validator in the committee.
                for (committee_i, validator_i) in cc.committee.iter().enumerate() {
                    // Assert the validators are assigned contiguously across committees.
                    assert_eq!(
                        *validator_i,
                        *expected_indices_iter.next().unwrap(),
                        "Non-sequential validators."
                    );
                    // Assert a call to `get_attestation_duties` is consistent with a call to
                    // `get_crosslink_committees_at_slot`
                    let attestation_duty = state
                        .get_attestation_duties(*validator_i, relative_epoch)
                        .unwrap()
                        .unwrap();
                    assert_eq!(attestation_duty.slot, slot);
                    assert_eq!(attestation_duty.shard, cc.shard);
                    assert_eq!(attestation_duty.committee_index, committee_i);
                    assert_eq!(attestation_duty.committee_len, cc.committee.len());
                }
            }
        }

        // Assert that all validators were assigned to a committee.
        assert!(expected_indices_iter.next().is_none());

        // Assert that all shards were assigned to a committee.
        assert!(expected_shards_iter.next().is_none());
    }

    fn committee_consistency_test<T: EthSpec>(
        validator_count: usize,
        state_epoch: Epoch,
        cache_epoch: RelativeEpoch,
    ) {
        let spec = &T::default_spec();

        let mut builder = TestingBeaconStateBuilder::from_single_keypair(
            validator_count,
            &Keypair::random(),
            spec,
        );

        let slot = state_epoch.start_slot(T::slots_per_epoch());
        builder.teleport_to_slot(slot);

        let (mut state, _keypairs): (BeaconState<T>, _) = builder.build();

        let distinct_hashes: Vec<Hash256> = (0..T::epochs_per_historical_vector())
            .into_iter()
            .map(|i| Hash256::from_low_u64_be(i as u64))
            .collect();
        state.randao_mixes = FixedVector::from(distinct_hashes);

        state
            .build_committee_cache(RelativeEpoch::Previous, spec)
            .unwrap();
        state
            .build_committee_cache(RelativeEpoch::Current, spec)
            .unwrap();
        state
            .build_committee_cache(RelativeEpoch::Next, spec)
            .unwrap();

        let cache_epoch = cache_epoch.into_epoch(state_epoch);

        execute_committee_consistency_test(state, cache_epoch, validator_count as usize, &spec);
    }

    fn committee_consistency_test_suite<T: EthSpec>(cached_epoch: RelativeEpoch) {
        let spec = T::default_spec();

        let validator_count = (T::shard_count() * spec.target_committee_size) + 1;

        committee_consistency_test::<T>(validator_count as usize, Epoch::new(0), cached_epoch);

        committee_consistency_test::<T>(
            validator_count as usize,
            T::genesis_epoch() + 4,
            cached_epoch,
        );

        committee_consistency_test::<T>(
            validator_count as usize,
            T::genesis_epoch() + T::slots_per_historical_root() as u64 * T::slots_per_epoch() * 4,
            cached_epoch,
        );
    }

    #[test]
    fn current_epoch_committee_consistency() {
        committee_consistency_test_suite::<MinimalEthSpec>(RelativeEpoch::Current);
    }

    #[test]
    fn previous_epoch_committee_consistency() {
        committee_consistency_test_suite::<MinimalEthSpec>(RelativeEpoch::Previous);
    }

    #[test]
    fn next_epoch_committee_consistency() {
        committee_consistency_test_suite::<MinimalEthSpec>(RelativeEpoch::Next);
    }
}
<|endoftext|>


<commit_before>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Martin Kysel <code@martinkysel.com>
//  *
//  * For the full copyright and license information, please view the LICENSE file
//  * that was distributed with this source code.

// spell-checker:ignore (methods) isnt

#[macro_use]
extern crate uucore;

use std::{
    fs::File,
    io::{stdin, stdout, BufReader, Read, Stdout, Write},
    path::Path,
    time::Duration,
};

#[cfg(all(unix, not(target_os = "fuchsia")))]
extern crate nix;

use clap::{crate_version, App, Arg};
use crossterm::{
    event::{self, Event, KeyCode, KeyEvent, KeyModifiers},
    execute, queue,
    style::Attribute,
    terminal,
};

use unicode_segmentation::UnicodeSegmentation;
use unicode_width::UnicodeWidthStr;
use uucore::display::Quotable;

const BELL: &str = "\x07";

pub mod options {
    pub const SILENT: &str = "silent";
    pub const LOGICAL: &str = "logical";
    pub const NO_PAUSE: &str = "no-pause";
    pub const PRINT_OVER: &str = "print-over";
    pub const CLEAN_PRINT: &str = "clean-print";
    pub const SQUEEZE: &str = "squeeze";
    pub const PLAIN: &str = "plain";
    pub const LINES: &str = "lines";
    pub const NUMBER: &str = "number";
    pub const PATTERN: &str = "pattern";
    pub const FROM_LINE: &str = "from-line";
    pub const FILES: &str = "files";
}

const MULTI_FILE_TOP_PROMPT: &str = "::::::::::::::\n{}\n::::::::::::::\n";

pub fn uumain(args: impl uucore::Args) -> i32 {
    let matches = uu_app().get_matches_from(args);

    let mut buff = String::new();
    let silent = matches.is_present(options::SILENT);
    if let Some(files) = matches.values_of(options::FILES) {
        let mut stdout = setup_term();
        let length = files.len();

        let mut files_iter = files.peekable();
        while let (Some(file), next_file) = (files_iter.next(), files_iter.peek()) {
            let file = Path::new(file);
            if file.is_dir() {
                terminal::disable_raw_mode().unwrap();
                show_usage_error!("{} is a directory.", file.quote());
                return 1;
            }
            if !file.exists() {
                terminal::disable_raw_mode().unwrap();
                show_error!("cannot open {}: No such file or directory", file.quote());
                return 1;
            }
            if length > 1 {
                buff.push_str(&MULTI_FILE_TOP_PROMPT.replace("{}", file.to_str().unwrap()));
            }
            let mut reader = BufReader::new(File::open(file).unwrap());
            reader.read_to_string(&mut buff).unwrap();
            more(&buff, &mut stdout, next_file.copied(), silent);
            buff.clear();
        }
        reset_term(&mut stdout);
    } else if atty::isnt(atty::Stream::Stdin) {
        stdin().read_to_string(&mut buff).unwrap();
        let mut stdout = setup_term();
        more(&buff, &mut stdout, None, silent);
        reset_term(&mut stdout);
    } else {
        show_usage_error!("bad usage");
    }
    0
}

pub fn uu_app() -> App<'static, 'static> {
    App::new(uucore::util_name())
        .about("A file perusal filter for CRT viewing.")
        .version(crate_version!())
        .arg(
            Arg::with_name(options::SILENT)
                .short("d")
                .long(options::SILENT)
                .help("Display help instead of ringing bell"),
        )
        // The commented arguments below are unimplemented:
        /*
        .arg(
            Arg::with_name(options::LOGICAL)
                .short("f")
                .long(options::LOGICAL)
                .help("Count logical rather than screen lines"),
        )
        .arg(
            Arg::with_name(options::NO_PAUSE)
                .short("l")
                .long(options::NO_PAUSE)
                .help("Suppress pause after form feed"),
        )
        .arg(
            Arg::with_name(options::PRINT_OVER)
                .short("c")
                .long(options::PRINT_OVER)
                .help("Do not scroll, display text and clean line ends"),
        )
        .arg(
            Arg::with_name(options::CLEAN_PRINT)
                .short("p")
                .long(options::CLEAN_PRINT)
                .help("Do not scroll, clean screen and display text"),
        )
        .arg(
            Arg::with_name(options::SQUEEZE)
                .short("s")
                .long(options::SQUEEZE)
                .help("Squeeze multiple blank lines into one"),
        )
        .arg(
            Arg::with_name(options::PLAIN)
                .short("u")
                .long(options::PLAIN)
                .help("Suppress underlining and bold"),
        )
        .arg(
            Arg::with_name(options::LINES)
                .short("n")
                .long(options::LINES)
                .value_name("number")
                .takes_value(true)
                .help("The number of lines per screen full"),
        )
        .arg(
            Arg::with_name(options::NUMBER)
                .allow_hyphen_values(true)
                .long(options::NUMBER)
                .required(false)
                .takes_value(true)
                .help("Same as --lines"),
        )
        .arg(
            Arg::with_name(options::FROM_LINE)
                .short("F")
                .allow_hyphen_values(true)
                .required(false)
                .takes_value(true)
                .value_name("number")
                .help("Display file beginning from line number"),
        )
        .arg(
            Arg::with_name(options::PATTERN)
                .short("P")
                .allow_hyphen_values(true)
                .required(false)
                .takes_value(true)
                .help("Display file beginning from pattern match"),
        )
        */
        .arg(
            Arg::with_name(options::FILES)
                .required(false)
                .multiple(true)
                .help("Path to the files to be read"),
        )
}

#[cfg(not(target_os = "fuchsia"))]
fn setup_term() -> std::io::Stdout {
    let stdout = stdout();
    terminal::enable_raw_mode().unwrap();
    stdout
}

#[cfg(target_os = "fuchsia")]
#[inline(always)]
fn setup_term() -> usize {
    0
}

#[cfg(not(target_os = "fuchsia"))]
fn reset_term(stdout: &mut std::io::Stdout) {
    terminal::disable_raw_mode().unwrap();
    // Clear the prompt
    queue!(stdout, terminal::Clear(terminal::ClearType::CurrentLine)).unwrap();
    // Move cursor to the beginning without printing new line
    print!("\r");
    stdout.flush().unwrap();
}

#[cfg(target_os = "fuchsia")]
#[inline(always)]
fn reset_term(_: &mut usize) {}

fn more(buff: &str, mut stdout: &mut Stdout, next_file: Option<&str>, silent: bool) {
    let (cols, rows) = terminal::size().unwrap();
    let lines = break_buff(buff, usize::from(cols));

    let mut pager = Pager::new(rows, lines, next_file, silent);
    pager.draw(stdout, None);
    if pager.should_close() {
        return;
    }

    loop {
        let mut wrong_key = None;
        if event::poll(Duration::from_millis(10)).unwrap() {
            match event::read().unwrap() {
                Event::Key(KeyEvent {
                    code: KeyCode::Char('q'),
                    modifiers: KeyModifiers::NONE,
                })
                | Event::Key(KeyEvent {
                    code: KeyCode::Char('c'),
                    modifiers: KeyModifiers::CONTROL,
                }) => {
                    reset_term(&mut stdout);
                    std::process::exit(0);
                }
                Event::Key(KeyEvent {
                    code: KeyCode::Down,
                    modifiers: KeyModifiers::NONE,
                })
                | Event::Key(KeyEvent {
                    code: KeyCode::Char(' '),
                    modifiers: KeyModifiers::NONE,
                }) => {
                    if pager.should_close() {
                        return;
                    } else {
                        pager.page_down();
                    }
                }
                Event::Key(KeyEvent {
                    code: KeyCode::Up,
                    modifiers: KeyModifiers::NONE,
                }) => {
                    pager.page_up();
                }
                Event::Resize(col, row) => {
                    pager.page_resize(col, row);
                }
                Event::Key(KeyEvent {
                    code: KeyCode::Char(k),
                    ..
                }) => wrong_key = Some(k),
                _ => continue,
            }

            pager.draw(stdout, wrong_key);
        }
    }
}

struct Pager<'a> {
    // The current line at the top of the screen
    upper_mark: usize,
    // The number of rows that fit on the screen
    content_rows: u16,
    lines: Vec<String>,
    next_file: Option<&'a str>,
    line_count: usize,
    silent: bool,
}

impl<'a> Pager<'a> {
    fn new(rows: u16, lines: Vec<String>, next_file: Option<&'a str>, silent: bool) -> Self {
        let line_count = lines.len();
        Self {
            upper_mark: 0,
            content_rows: rows.saturating_sub(1),
            lines,
            next_file,
            line_count,
            silent,
        }
    }

    fn should_close(&mut self) -> bool {
        self.upper_mark
            .saturating_add(self.content_rows.into())
            .ge(&self.line_count)
    }

    fn page_down(&mut self) {
        self.upper_mark = self.upper_mark.saturating_add(self.content_rows.into());
    }

    fn page_up(&mut self) {
        self.upper_mark = self.upper_mark.saturating_sub(self.content_rows.into());
    }

    // TODO: Deal with column size changes.
    fn page_resize(&mut self, _: u16, row: u16) {
        self.content_rows = row.saturating_sub(1);
    }

    fn draw(&self, stdout: &mut std::io::Stdout, wrong_key: Option<char>) {
        let lower_mark = self
            .line_count
            .min(self.upper_mark.saturating_add(self.content_rows.into()));
        self.draw_lines(stdout);
        self.draw_prompt(stdout, lower_mark, wrong_key);
        stdout.flush().unwrap();
    }

    fn draw_lines(&self, stdout: &mut std::io::Stdout) {
        execute!(stdout, terminal::Clear(terminal::ClearType::CurrentLine)).unwrap();
        let displayed_lines = self
            .lines
            .iter()
            .skip(self.upper_mark)
            .take(self.content_rows.into());

        for line in displayed_lines {
            stdout
                .write_all(format!("\r{}\n", line).as_bytes())
                .unwrap();
        }
    }

    fn draw_prompt(&self, stdout: &mut Stdout, lower_mark: usize, wrong_key: Option<char>) {
        let status_inner = if lower_mark == self.line_count {
            format!("Next file: {}", self.next_file.unwrap_or_default())
        } else {
            format!(
                "{}%",
                (lower_mark as f64 / self.line_count as f64 * 100.0).round() as u16
            )
        };

        let status = format!("--More--({})", status_inner);

        let banner = match (self.silent, wrong_key) {
            (true, Some(key)) => format!(
                "{} [Unknown key: '{}'. Press 'h' for instructions. (unimplemented)]",
                status, key
            ),
            (true, None) => format!("{}[Press space to continue, 'q' to quit.]", status),
            (false, Some(_)) => format!("{}{}", status, BELL),
            (false, None) => status,
        };

        write!(
            stdout,
            "\r{}{}{}",
            Attribute::Reverse,
            banner,
            Attribute::Reset
        )
        .unwrap();
    }
}

// Break the lines on the cols of the terminal
fn break_buff(buff: &str, cols: usize) -> Vec<String> {
    let mut lines = Vec::with_capacity(buff.lines().count());

    for l in buff.lines() {
        lines.append(&mut break_line(l, cols));
    }
    lines
}

fn break_line(line: &str, cols: usize) -> Vec<String> {
    let width = UnicodeWidthStr::width(line);
    let mut lines = Vec::new();
    if width < cols {
        lines.push(line.to_string());
        return lines;
    }

    let gr_idx = UnicodeSegmentation::grapheme_indices(line, true);
    let mut last_index = 0;
    let mut total_width = 0;
    for (index, grapheme) in gr_idx {
        let width = UnicodeWidthStr::width(grapheme);
        total_width += width;

        if total_width > cols {
            lines.push(line[last_index..index].to_string());
            last_index = index;
            total_width = width;
        }
    }

    if last_index != line.len() {
        lines.push(line[last_index..].to_string());
    }
    lines
}

#[cfg(test)]
mod tests {
    use super::break_line;
    use unicode_width::UnicodeWidthStr;

    #[test]
    fn test_break_lines_long() {
        let mut test_string = String::with_capacity(100);
        for _ in 0..200 {
            test_string.push('#');
        }

        let lines = break_line(&test_string, 80);
        let widths: Vec<usize> = lines
            .iter()
            .map(|s| UnicodeWidthStr::width(&s[..]))
            .collect();

        assert_eq!((80, 80, 40), (widths[0], widths[1], widths[2]));
    }

    #[test]
    fn test_break_lines_short() {
        let mut test_string = String::with_capacity(100);
        for _ in 0..20 {
            test_string.push('#');
        }

        let lines = break_line(&test_string, 80);

        assert_eq!(20, lines[0].len());
    }

    #[test]
    fn test_break_line_zwj() {
        let mut test_string = String::with_capacity(1100);
        for _ in 0..20 {
            test_string.push_str("👩🏻🔬");
        }

        let lines = break_line(&test_string, 80);

        let widths: Vec<usize> = lines
            .iter()
            .map(|s| UnicodeWidthStr::width(&s[..]))
            .collect();

        // Each 👩🏻🔬 is 6 character width it break line to the closest number to 80 => 6 * 13 = 78
        assert_eq!((78, 42), (widths[0], widths[1]));
    }
}
<commit_msg>more: remove redundant mut from stdout accesses<commit_after>//  * This file is part of the uutils coreutils package.
//  *
//  * (c) Martin Kysel <code@martinkysel.com>
//  *
//  * For the full copyright and license information, please view the LICENSE file
//  * that was distributed with this source code.

// spell-checker:ignore (methods) isnt

#[macro_use]
extern crate uucore;

use std::{
    fs::File,
    io::{stdin, stdout, BufReader, Read, Stdout, Write},
    path::Path,
    time::Duration,
};

#[cfg(all(unix, not(target_os = "fuchsia")))]
extern crate nix;

use clap::{crate_version, App, Arg};
use crossterm::{
    event::{self, Event, KeyCode, KeyEvent, KeyModifiers},
    execute, queue,
    style::Attribute,
    terminal,
};

use unicode_segmentation::UnicodeSegmentation;
use unicode_width::UnicodeWidthStr;
use uucore::display::Quotable;

const BELL: &str = "\x07";

pub mod options {
    pub const SILENT: &str = "silent";
    pub const LOGICAL: &str = "logical";
    pub const NO_PAUSE: &str = "no-pause";
    pub const PRINT_OVER: &str = "print-over";
    pub const CLEAN_PRINT: &str = "clean-print";
    pub const SQUEEZE: &str = "squeeze";
    pub const PLAIN: &str = "plain";
    pub const LINES: &str = "lines";
    pub const NUMBER: &str = "number";
    pub const PATTERN: &str = "pattern";
    pub const FROM_LINE: &str = "from-line";
    pub const FILES: &str = "files";
}

const MULTI_FILE_TOP_PROMPT: &str = "::::::::::::::\n{}\n::::::::::::::\n";

pub fn uumain(args: impl uucore::Args) -> i32 {
    let matches = uu_app().get_matches_from(args);

    let mut buff = String::new();
    let silent = matches.is_present(options::SILENT);
    if let Some(files) = matches.values_of(options::FILES) {
        let mut stdout = setup_term();
        let length = files.len();

        let mut files_iter = files.peekable();
        while let (Some(file), next_file) = (files_iter.next(), files_iter.peek()) {
            let file = Path::new(file);
            if file.is_dir() {
                terminal::disable_raw_mode().unwrap();
                show_usage_error!("{} is a directory.", file.quote());
                return 1;
            }
            if !file.exists() {
                terminal::disable_raw_mode().unwrap();
                show_error!("cannot open {}: No such file or directory", file.quote());
                return 1;
            }
            if length > 1 {
                buff.push_str(&MULTI_FILE_TOP_PROMPT.replace("{}", file.to_str().unwrap()));
            }
            let mut reader = BufReader::new(File::open(file).unwrap());
            reader.read_to_string(&mut buff).unwrap();
            more(&buff, &mut stdout, next_file.copied(), silent);
            buff.clear();
        }
        reset_term(&mut stdout);
    } else if atty::isnt(atty::Stream::Stdin) {
        stdin().read_to_string(&mut buff).unwrap();
        let mut stdout = setup_term();
        more(&buff, &mut stdout, None, silent);
        reset_term(&mut stdout);
    } else {
        show_usage_error!("bad usage");
    }
    0
}

pub fn uu_app() -> App<'static, 'static> {
    App::new(uucore::util_name())
        .about("A file perusal filter for CRT viewing.")
        .version(crate_version!())
        .arg(
            Arg::with_name(options::SILENT)
                .short("d")
                .long(options::SILENT)
                .help("Display help instead of ringing bell"),
        )
        // The commented arguments below are unimplemented:
        /*
        .arg(
            Arg::with_name(options::LOGICAL)
                .short("f")
                .long(options::LOGICAL)
                .help("Count logical rather than screen lines"),
        )
        .arg(
            Arg::with_name(options::NO_PAUSE)
                .short("l")
                .long(options::NO_PAUSE)
                .help("Suppress pause after form feed"),
        )
        .arg(
            Arg::with_name(options::PRINT_OVER)
                .short("c")
                .long(options::PRINT_OVER)
                .help("Do not scroll, display text and clean line ends"),
        )
        .arg(
            Arg::with_name(options::CLEAN_PRINT)
                .short("p")
                .long(options::CLEAN_PRINT)
                .help("Do not scroll, clean screen and display text"),
        )
        .arg(
            Arg::with_name(options::SQUEEZE)
                .short("s")
                .long(options::SQUEEZE)
                .help("Squeeze multiple blank lines into one"),
        )
        .arg(
            Arg::with_name(options::PLAIN)
                .short("u")
                .long(options::PLAIN)
                .help("Suppress underlining and bold"),
        )
        .arg(
            Arg::with_name(options::LINES)
                .short("n")
                .long(options::LINES)
                .value_name("number")
                .takes_value(true)
                .help("The number of lines per screen full"),
        )
        .arg(
            Arg::with_name(options::NUMBER)
                .allow_hyphen_values(true)
                .long(options::NUMBER)
                .required(false)
                .takes_value(true)
                .help("Same as --lines"),
        )
        .arg(
            Arg::with_name(options::FROM_LINE)
                .short("F")
                .allow_hyphen_values(true)
                .required(false)
                .takes_value(true)
                .value_name("number")
                .help("Display file beginning from line number"),
        )
        .arg(
            Arg::with_name(options::PATTERN)
                .short("P")
                .allow_hyphen_values(true)
                .required(false)
                .takes_value(true)
                .help("Display file beginning from pattern match"),
        )
        */
        .arg(
            Arg::with_name(options::FILES)
                .required(false)
                .multiple(true)
                .help("Path to the files to be read"),
        )
}

#[cfg(not(target_os = "fuchsia"))]
fn setup_term() -> std::io::Stdout {
    let stdout = stdout();
    terminal::enable_raw_mode().unwrap();
    stdout
}

#[cfg(target_os = "fuchsia")]
#[inline(always)]
fn setup_term() -> usize {
    0
}

#[cfg(not(target_os = "fuchsia"))]
fn reset_term(stdout: &mut std::io::Stdout) {
    terminal::disable_raw_mode().unwrap();
    // Clear the prompt
    queue!(stdout, terminal::Clear(terminal::ClearType::CurrentLine)).unwrap();
    // Move cursor to the beginning without printing new line
    print!("\r");
    stdout.flush().unwrap();
}

#[cfg(target_os = "fuchsia")]
#[inline(always)]
fn reset_term(_: &mut usize) {}

fn more(buff: &str, stdout: &mut Stdout, next_file: Option<&str>, silent: bool) {
    let (cols, rows) = terminal::size().unwrap();
    let lines = break_buff(buff, usize::from(cols));

    let mut pager = Pager::new(rows, lines, next_file, silent);
    pager.draw(stdout, None);
    if pager.should_close() {
        return;
    }

    loop {
        let mut wrong_key = None;
        if event::poll(Duration::from_millis(10)).unwrap() {
            match event::read().unwrap() {
                Event::Key(KeyEvent {
                    code: KeyCode::Char('q'),
                    modifiers: KeyModifiers::NONE,
                })
                | Event::Key(KeyEvent {
                    code: KeyCode::Char('c'),
                    modifiers: KeyModifiers::CONTROL,
                }) => {
                    reset_term(stdout);
                    std::process::exit(0);
                }
                Event::Key(KeyEvent {
                    code: KeyCode::Down,
                    modifiers: KeyModifiers::NONE,
                })
                | Event::Key(KeyEvent {
                    code: KeyCode::Char(' '),
                    modifiers: KeyModifiers::NONE,
                }) => {
                    if pager.should_close() {
                        return;
                    } else {
                        pager.page_down();
                    }
                }
                Event::Key(KeyEvent {
                    code: KeyCode::Up,
                    modifiers: KeyModifiers::NONE,
                }) => {
                    pager.page_up();
                }
                Event::Resize(col, row) => {
                    pager.page_resize(col, row);
                }
                Event::Key(KeyEvent {
                    code: KeyCode::Char(k),
                    ..
                }) => wrong_key = Some(k),
                _ => continue,
            }

            pager.draw(stdout, wrong_key);
        }
    }
}

struct Pager<'a> {
    // The current line at the top of the screen
    upper_mark: usize,
    // The number of rows that fit on the screen
    content_rows: u16,
    lines: Vec<String>,
    next_file: Option<&'a str>,
    line_count: usize,
    silent: bool,
}

impl<'a> Pager<'a> {
    fn new(rows: u16, lines: Vec<String>, next_file: Option<&'a str>, silent: bool) -> Self {
        let line_count = lines.len();
        Self {
            upper_mark: 0,
            content_rows: rows.saturating_sub(1),
            lines,
            next_file,
            line_count,
            silent,
        }
    }

    fn should_close(&mut self) -> bool {
        self.upper_mark
            .saturating_add(self.content_rows.into())
            .ge(&self.line_count)
    }

    fn page_down(&mut self) {
        self.upper_mark = self.upper_mark.saturating_add(self.content_rows.into());
    }

    fn page_up(&mut self) {
        self.upper_mark = self.upper_mark.saturating_sub(self.content_rows.into());
    }

    // TODO: Deal with column size changes.
    fn page_resize(&mut self, _: u16, row: u16) {
        self.content_rows = row.saturating_sub(1);
    }

    fn draw(&self, stdout: &mut std::io::Stdout, wrong_key: Option<char>) {
        let lower_mark = self
            .line_count
            .min(self.upper_mark.saturating_add(self.content_rows.into()));
        self.draw_lines(stdout);
        self.draw_prompt(stdout, lower_mark, wrong_key);
        stdout.flush().unwrap();
    }

    fn draw_lines(&self, stdout: &mut std::io::Stdout) {
        execute!(stdout, terminal::Clear(terminal::ClearType::CurrentLine)).unwrap();
        let displayed_lines = self
            .lines
            .iter()
            .skip(self.upper_mark)
            .take(self.content_rows.into());

        for line in displayed_lines {
            stdout
                .write_all(format!("\r{}\n", line).as_bytes())
                .unwrap();
        }
    }

    fn draw_prompt(&self, stdout: &mut Stdout, lower_mark: usize, wrong_key: Option<char>) {
        let status_inner = if lower_mark == self.line_count {
            format!("Next file: {}", self.next_file.unwrap_or_default())
        } else {
            format!(
                "{}%",
                (lower_mark as f64 / self.line_count as f64 * 100.0).round() as u16
            )
        };

        let status = format!("--More--({})", status_inner);

        let banner = match (self.silent, wrong_key) {
            (true, Some(key)) => format!(
                "{} [Unknown key: '{}'. Press 'h' for instructions. (unimplemented)]",
                status, key
            ),
            (true, None) => format!("{}[Press space to continue, 'q' to quit.]", status),
            (false, Some(_)) => format!("{}{}", status, BELL),
            (false, None) => status,
        };

        write!(
            stdout,
            "\r{}{}{}",
            Attribute::Reverse,
            banner,
            Attribute::Reset
        )
        .unwrap();
    }
}

// Break the lines on the cols of the terminal
fn break_buff(buff: &str, cols: usize) -> Vec<String> {
    let mut lines = Vec::with_capacity(buff.lines().count());

    for l in buff.lines() {
        lines.append(&mut break_line(l, cols));
    }
    lines
}

fn break_line(line: &str, cols: usize) -> Vec<String> {
    let width = UnicodeWidthStr::width(line);
    let mut lines = Vec::new();
    if width < cols {
        lines.push(line.to_string());
        return lines;
    }

    let gr_idx = UnicodeSegmentation::grapheme_indices(line, true);
    let mut last_index = 0;
    let mut total_width = 0;
    for (index, grapheme) in gr_idx {
        let width = UnicodeWidthStr::width(grapheme);
        total_width += width;

        if total_width > cols {
            lines.push(line[last_index..index].to_string());
            last_index = index;
            total_width = width;
        }
    }

    if last_index != line.len() {
        lines.push(line[last_index..].to_string());
    }
    lines
}

#[cfg(test)]
mod tests {
    use super::break_line;
    use unicode_width::UnicodeWidthStr;

    #[test]
    fn test_break_lines_long() {
        let mut test_string = String::with_capacity(100);
        for _ in 0..200 {
            test_string.push('#');
        }

        let lines = break_line(&test_string, 80);
        let widths: Vec<usize> = lines
            .iter()
            .map(|s| UnicodeWidthStr::width(&s[..]))
            .collect();

        assert_eq!((80, 80, 40), (widths[0], widths[1], widths[2]));
    }

    #[test]
    fn test_break_lines_short() {
        let mut test_string = String::with_capacity(100);
        for _ in 0..20 {
            test_string.push('#');
        }

        let lines = break_line(&test_string, 80);

        assert_eq!(20, lines[0].len());
    }

    #[test]
    fn test_break_line_zwj() {
        let mut test_string = String::with_capacity(1100);
        for _ in 0..20 {
            test_string.push_str("👩🏻🔬");
        }

        let lines = break_line(&test_string, 80);

        let widths: Vec<usize> = lines
            .iter()
            .map(|s| UnicodeWidthStr::width(&s[..]))
            .collect();

        // Each 👩🏻🔬 is 6 character width it break line to the closest number to 80 => 6 * 13 = 78
        assert_eq!((78, 42), (widths[0], widths[1]));
    }
}
<|endoftext|>


<commit_before>//! Helper to crudely rewrite types from the spec's format to our own.
//!
//! This reads a file containing table and record descriptions (from the spec)
//! and converts them to the form that we use, writing the result to stdout.
//!
//! For more information about how this works, see the README

use std::{fmt::Write, ops::Deref};

macro_rules! exit_with_msg {
    ($disp:expr, $line:expr) => {{
        eprintln!("ERROR: {}", $disp);
        eprintln!("Line {}: '{}'", 1 + $line.number, $line.text);
        std::process::exit(1);
    }};
}

/// a wrapper around a line, so we can report errors with line numbers
struct Line<'a> {
    text: &'a str,
    number: usize,
}

impl Deref for Line<'_> {
    type Target = str;
    fn deref(&self) -> &<Self as Deref>::Target {
        self.text
    }
}

fn main() {
    let in_path = std::env::args().nth(1).expect("expected path argument");
    let input = std::fs::read_to_string(in_path).expect("failed to read path");
    let mut lines = input
        .lines()
        .enumerate()
        .map(|(number, text)| Line {
            text: text.trim(),
            number,
        })
        .filter(|l| !l.starts_with('#'));

    while let Some(item) = generate_one_item(&mut lines) {
        println!("{}\n", item)
    }
}

/// parse a single enum, table or record.
///
/// Returns `Some` on success, `None` if there are no more items, and terminates
/// if something goes wrong.
fn generate_one_item<'a>(lines: impl Iterator<Item = Line<'a>>) -> Option<String> {
    let mut lines = lines.skip_while(|line| line.is_empty());
    let mut comments = Vec::new();
    let decl = loop {
        match lines.next() {
            Some(line) if line.starts_with("///") => comments.push(line.text),
            Some(line) if line.starts_with('@') => break Decl::parse(line).unwrap(),
            Some(line) => exit_with_msg!("expected table or record name", line),
            None => return None,
        }
    };

    let item = match decl.kind {
        DeclKind::RawEnum => generate_one_enum(decl, lines),
        DeclKind::Table | DeclKind::Record => generate_one_table(decl, lines),
        DeclKind::Flags => generate_one_flags(decl, lines),
    }?;
    let mut comments = comments.join("\n");
    if comments.is_empty() {
        Some(item)
    } else {
        comments.push('\n');
        comments.push_str(&item);
        Some(comments)
    }
}

/// Generate a single table or record (they're currently the same)
fn generate_one_table<'a>(decl: Decl, lines: impl Iterator<Item = Line<'a>>) -> Option<String> {
    let fields = lines.map_while(parse_field).collect::<Vec<_>>();
    let lifetime_str = if fields.iter().any(|x| x.maybe_count.is_some()) {
        "<'a>"
    } else {
        ""
    };
    let mut result = String::new();
    writeln!(&mut result, "{}{} {{", decl.name, lifetime_str).unwrap();
    for line in &fields {
        writeln!(&mut result, "{}", line).unwrap();
    }
    result.push('}');
    Some(result)
}

fn generate_one_enum<'a>(decl: Decl, lines: impl Iterator<Item = Line<'a>>) -> Option<String> {
    let fields = lines.map_while(parse_field).collect::<Vec<_>>();
    let mut result = String::new();
    writeln!(
        &mut result,
        "#[repr({})]\nenum {} {{",
        decl.annotation, decl.name
    )
    .unwrap();
    for line in &fields {
        writeln!(&mut result, "    {} = {},", line.name, line.typ).unwrap();
    }
    result.push('}');
    Some(result)
}

fn generate_one_flags<'a>(decl: Decl, lines: impl Iterator<Item = Line<'a>>) -> Option<String> {
    let fields = lines.map_while(parse_field).collect::<Vec<_>>();
    let mut result = String::new();
    writeln!(
        &mut result,
        "#[flags({})]\n{} {{",
        decl.annotation, decl.name
    )
    .unwrap();
    for line in &fields {
        format_comment(&mut result, "    ", line.comment).unwrap();
        writeln!(&mut result, "    {} = {},", line.name, line.typ).unwrap();
    }
    result.push('}');
    Some(result)
}

enum DeclKind {
    Table,
    Record,
    RawEnum,
    Flags,
}

struct Decl<'a> {
    kind: DeclKind,
    annotation: &'a str,
    name: &'a str,
}

impl<'a> Decl<'a> {
    fn parse(line: Line<'a>) -> Option<Self> {
        let mut decl = line.text.split_whitespace();
        let mut annotation = "";
        let kind = match decl.next()? {
            "@table" => DeclKind::Table,
            "@record" => DeclKind::Record,
            x if x.starts_with("@enum(") || x.starts_with("@flags(") => {
                let repr = x.split_once('(').unwrap().1.trim_end_matches(')');
                //let repr = x.trim_start_matches("@enum(").trim_end_matches(')');
                if !["u8", "u16"].contains(&repr) {
                    exit_with_msg!(format!("unexpected enum/flag repr '{}'", repr), line);
                }
                annotation = repr;
                if x.starts_with("@enum") {
                    DeclKind::RawEnum
                } else {
                    DeclKind::Flags
                }
            }
            "@enum" | "@flags" => exit_with_msg!(
                "@enum/@flags requires explicit repr like: '@flags(u16)'",
                line
            ),
            other => exit_with_msg!(format!("unknown item kind '{}'", other), line),
        };
        let name = decl
            .next()
            .unwrap_or_else(|| exit_with_msg!("missing name", line));

        Some(Decl {
            kind,
            annotation,
            name,
        })
    }
}

struct Field<'a> {
    name: &'a str,
    maybe_count: Option<String>,
    typ: &'a str,
    comment: &'a str,
}

impl<'a> std::fmt::Display for Field<'a> {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        format_comment(f, "    ", self.comment)?;
        if self.name.contains("reserved") {
            writeln!(f, "    #[hidden]")?;
        }
        if let Some(count) = &self.maybe_count {
            writeln!(f, "    #[count({})]", count)?;
            write!(f, "    {}: [{}],", decamalize(self.name), self.typ)?;
        } else {
            write!(f, "    {}: {},", decamalize(self.name), self.typ)?;
        }
        Ok(())
    }
}

fn parse_field(line: Line) -> Option<Field> {
    if line.is_empty() {
        return None;
    }
    let mut iter = line.text.splitn(3, '\t');
    let (typ, ident, comment) = match (iter.next(), iter.next(), iter.next()) {
        (Some(a), Some(b), Some(c)) => (a, b, c),
        (Some(a), Some(b), None) => (a, b, ""),
        _ => exit_with_msg!("line could not be parsed as type/name/comment", line),
    };
    let typ = normalize_type(typ);
    let (name, maybe_count) = split_ident(ident);
    //let name = decamalize(name);
    let maybe_count = maybe_count.map(decamalize);
    Some(Field {
        name,
        maybe_count,
        typ,
        comment,
    })
}

/// takes an ident and splits it into the name and an optional count (if the item
/// is an array)
fn split_ident(input: &str) -> (&str, Option<&str>) {
    match input.split_once('[') {
        Some((front, back)) => (front, Some(back.trim_end_matches(']'))),
        None => (input, None),
    }
}

fn normalize_type(input: &str) -> &str {
    match input {
        "uint8" => "BigEndian<u8>",
        "uint16" => "BigEndian<u16>",
        "uint24" => "BigEndian<Uint24>",
        "uint32" => "BigEndian<u32>",
        "int8" => "BigEndian<i8>",
        "int16" => "BigEndian<i16>",
        "int32" => "BigEndian<i32>",
        "FWORD" => "BigEndian<FWord>",
        "UFWORD" => "BigEndian<UfWord>",
        "F2DOT14" => "BigEndian<F2Dot14>",
        "LONGDATETIME" => "BigEndian<LongDateTime>",
        "Version16Dot16" => "BigEndian<Version16Dot16>",
        "Fixed" => "BigEndian<Fixed>",
        "Tag" => "BigEndian<Tag>",
        "Offset16" => "BigEndian<Offset16>",
        "Offset24" => "BigEndian<Offset24>",
        "Offset32" => "BigEndian<Offset32>",
        other => other,
    }
}

fn decamalize(input: &str) -> String {
    //taken from serde: https://github.com/serde-rs/serde/blob/7e19ae8c9486a3bbbe51f1befb05edee94c454f9/serde_derive/src/internals/case.rs#L69-L76
    let mut snake = String::new();
    for (i, ch) in input.char_indices() {
        if i > 0 && ch.is_uppercase() {
            snake.push('_');
        }
        snake.push(ch.to_ascii_lowercase());
    }
    snake
}

fn format_comment(f: &mut dyn std::fmt::Write, whitespace: &str, input: &str) -> std::fmt::Result {
    const LINE_LEN: usize = 72;

    let mut cur_len = 0;

    for token in input.split_inclusive(' ') {
        if cur_len == 0 || cur_len + token.len() > LINE_LEN {
            if cur_len > 0 {
                writeln!(f)?;
            }
            write!(f, "{}/// ", whitespace)?;
            cur_len = whitespace.len() + 4;
        }
        write!(f, "{}", token)?;
        cur_len += token.len();
    }
    if cur_len > 0 {
        writeln!(f)?;
    }
    Ok(())
}
<commit_msg>[codegen] Update preprocessor<commit_after>//! Helper to crudely rewrite types from the spec's format to our own.
//!
//! This reads a file containing table and record descriptions (from the spec)
//! and converts them to the form that we use, writing the result to stdout.
//!
//! For more information about how this works, see the README

use std::{fmt::Write, ops::Deref};

macro_rules! exit_with_msg {
    ($disp:expr, $line:expr) => {{
        eprintln!("ERROR: {}", $disp);
        eprintln!("Line {}: '{}'", 1 + $line.number, $line.text);
        std::process::exit(1);
    }};
}

/// a wrapper around a line, so we can report errors with line numbers
struct Line<'a> {
    text: &'a str,
    number: usize,
}

impl Deref for Line<'_> {
    type Target = str;
    fn deref(&self) -> &<Self as Deref>::Target {
        self.text
    }
}

fn main() {
    let in_path = std::env::args().nth(1).expect("expected path argument");
    let input = std::fs::read_to_string(in_path).expect("failed to read path");
    let mut lines = input
        .lines()
        .enumerate()
        .map(|(number, text)| Line {
            text: text.trim(),
            number,
        })
        .filter(|l| !l.starts_with('#'));

    while let Some(item) = generate_one_item(&mut lines) {
        println!("{}\n", item)
    }
}

/// parse a single enum, table or record.
///
/// Returns `Some` on success, `None` if there are no more items, and terminates
/// if something goes wrong.
fn generate_one_item<'a>(lines: impl Iterator<Item = Line<'a>>) -> Option<String> {
    let mut lines = lines.skip_while(|line| line.is_empty());
    let mut comments = Vec::new();
    let decl = loop {
        match lines.next() {
            Some(line) if line.starts_with("///") => comments.push(line.text),
            Some(line) if line.starts_with('@') => break Decl::parse(line).unwrap(),
            Some(line) => exit_with_msg!("expected table or record name", line),
            None => return None,
        }
    };

    let item = match decl.kind {
        DeclKind::RawEnum => generate_one_enum(decl, lines),
        DeclKind::Table | DeclKind::Record => generate_one_table(decl, lines),
        DeclKind::Flags => generate_one_flags(decl, lines),
    }?;
    let mut comments = comments.join("\n");
    if comments.is_empty() {
        Some(item)
    } else {
        comments.push('\n');
        comments.push_str(&item);
        Some(comments)
    }
}

/// Generate a single table or record (they're currently the same)
fn generate_one_table<'a>(decl: Decl, lines: impl Iterator<Item = Line<'a>>) -> Option<String> {
    let fields = lines.map_while(parse_field).collect::<Vec<_>>();
    let mut result = String::new();
    writeln!(&mut result, "{} {} {{", decl.kind, decl.name).unwrap();
    for line in &fields {
        writeln!(&mut result, "{}", line).unwrap();
    }
    result.push('}');
    Some(result)
}

fn generate_one_enum<'a>(decl: Decl, lines: impl Iterator<Item = Line<'a>>) -> Option<String> {
    let fields = lines.map_while(parse_field).collect::<Vec<_>>();
    let mut result = String::new();
    writeln!(
        &mut result,
        "#[repr({})]\nenum {} {{",
        decl.annotation, decl.name
    )
    .unwrap();
    for line in &fields {
        writeln!(&mut result, "    {} = {},", line.name, line.typ).unwrap();
    }
    result.push('}');
    Some(result)
}

fn generate_one_flags<'a>(decl: Decl, lines: impl Iterator<Item = Line<'a>>) -> Option<String> {
    let fields = lines.map_while(parse_field).collect::<Vec<_>>();
    let mut result = String::new();
    writeln!(
        &mut result,
        "#[flags({})]\n{} {{",
        decl.annotation, decl.name
    )
    .unwrap();
    for line in &fields {
        format_comment(&mut result, "    ", line.comment).unwrap();
        writeln!(&mut result, "    {} = {},", line.name, line.typ).unwrap();
    }
    result.push('}');
    Some(result)
}

enum DeclKind {
    Table,
    Record,
    RawEnum,
    Flags,
}

impl std::fmt::Display for DeclKind {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
            DeclKind::Table => write!(f, "table"),
            DeclKind::Record => write!(f, "record"),
            DeclKind::RawEnum => write!(f, "enum"),
            DeclKind::Flags => write!(f, "flags"),
        }
    }
}

struct Decl<'a> {
    kind: DeclKind,
    annotation: &'a str,
    name: &'a str,
}

impl<'a> Decl<'a> {
    fn parse(line: Line<'a>) -> Option<Self> {
        let mut decl = line.text.split_whitespace();
        let mut annotation = "";
        let kind = match decl.next()? {
            "@table" => DeclKind::Table,
            "@record" => DeclKind::Record,
            x if x.starts_with("@enum(") || x.starts_with("@flags(") => {
                let repr = x.split_once('(').unwrap().1.trim_end_matches(')');
                //let repr = x.trim_start_matches("@enum(").trim_end_matches(')');
                if !["u8", "u16"].contains(&repr) {
                    exit_with_msg!(format!("unexpected enum/flag repr '{}'", repr), line);
                }
                annotation = repr;
                if x.starts_with("@enum") {
                    DeclKind::RawEnum
                } else {
                    DeclKind::Flags
                }
            }
            "@enum" | "@flags" => exit_with_msg!(
                "@enum/@flags requires explicit repr like: '@flags(u16)'",
                line
            ),
            other => exit_with_msg!(format!("unknown item kind '{}'", other), line),
        };
        let name = decl
            .next()
            .unwrap_or_else(|| exit_with_msg!("missing name", line));

        Some(Decl {
            kind,
            annotation,
            name,
        })
    }
}

struct Field<'a> {
    name: &'a str,
    maybe_count: Option<String>,
    typ: &'a str,
    comment: &'a str,
}

impl<'a> std::fmt::Display for Field<'a> {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        format_comment(f, "    ", self.comment)?;
        if self.name.contains("reserved") {
            writeln!(f, "    #[hidden]")?;
        }
        if let Some(count) = &self.maybe_count {
            writeln!(f, "    #[count(${})]", count)?;
            write!(f, "    {}: [{}],", decamalize(self.name), self.typ)?;
        } else {
            write!(f, "    {}: {},", decamalize(self.name), self.typ)?;
        }
        Ok(())
    }
}

fn parse_field(line: Line) -> Option<Field> {
    if line.is_empty() {
        return None;
    }
    let mut iter = line.text.splitn(3, '\t');
    let (typ, ident, comment) = match (iter.next(), iter.next(), iter.next()) {
        (Some(a), Some(b), Some(c)) => (a, b, c),
        (Some(a), Some(b), None) => (a, b, ""),
        _ => exit_with_msg!("line could not be parsed as type/name/comment", line),
    };
    let typ = normalize_type(typ);
    let (name, maybe_count) = split_ident(ident);
    //let name = decamalize(name);
    let maybe_count = maybe_count.map(decamalize);
    Some(Field {
        name,
        maybe_count,
        typ,
        comment,
    })
}

/// takes an ident and splits it into the name and an optional count (if the item
/// is an array)
fn split_ident(input: &str) -> (&str, Option<&str>) {
    match input.split_once('[') {
        Some((front, back)) => (front, Some(back.trim_end_matches(']'))),
        None => (input, None),
    }
}

fn normalize_type(input: &str) -> &str {
    match input {
        "uint8" => "BigEndian<u8>",
        "uint16" => "BigEndian<u16>",
        "uint24" => "BigEndian<Uint24>",
        "uint32" => "BigEndian<u32>",
        "int8" => "BigEndian<i8>",
        "int16" => "BigEndian<i16>",
        "int32" => "BigEndian<i32>",
        "FWORD" => "BigEndian<FWord>",
        "UFWORD" => "BigEndian<UfWord>",
        "F2DOT14" => "BigEndian<F2Dot14>",
        "LONGDATETIME" => "BigEndian<LongDateTime>",
        "Version16Dot16" => "BigEndian<Version16Dot16>",
        "Fixed" => "BigEndian<Fixed>",
        "Tag" => "BigEndian<Tag>",
        "Offset16" => "BigEndian<Offset16>",
        "Offset24" => "BigEndian<Offset24>",
        "Offset32" => "BigEndian<Offset32>",
        other => other,
    }
}

fn decamalize(input: &str) -> String {
    //taken from serde: https://github.com/serde-rs/serde/blob/7e19ae8c9486a3bbbe51f1befb05edee94c454f9/serde_derive/src/internals/case.rs#L69-L76
    let mut snake = String::new();
    for (i, ch) in input.char_indices() {
        if i > 0 && ch.is_uppercase() {
            snake.push('_');
        }
        snake.push(ch.to_ascii_lowercase());
    }
    snake
}

fn format_comment(f: &mut dyn std::fmt::Write, whitespace: &str, input: &str) -> std::fmt::Result {
    const LINE_LEN: usize = 72;

    let mut cur_len = 0;

    for token in input.split_inclusive(' ') {
        if cur_len == 0 || cur_len + token.len() > LINE_LEN {
            if cur_len > 0 {
                writeln!(f)?;
            }
            write!(f, "{}/// ", whitespace)?;
            cur_len = whitespace.len() + 4;
        }
        write!(f, "{}", token)?;
        cur_len += token.len();
    }
    if cur_len > 0 {
        writeln!(f)?;
    }
    Ok(())
}
<|endoftext|>


<commit_before>extern crate ktensor as k;
use k::{Arc, Vec2, Tensor, Context, Graph, State, Variable};
use std::path::{Path};
use std::fs::{File};
use std::io::{BufReader, Read};
use std::error::{Error};
use std::cmp;

fn read_u8(reader: &mut Read) -> u8 {
    use std::mem;

    let mut buf: [u8; 1] = [0];
    match reader.read_exact(&mut buf).map(|_| {
        let data: u8 = unsafe { mem::transmute(buf) };
        data
    }) {
        Err(reason) => panic!("failed to read u8: {}", Error::description(&reason)),
        Ok(byte)    => byte,
    }
}

fn read_u32(reader: &mut Read) -> u32 {
    use std::mem;

    let mut buf: [u8; 4] = [0, 0, 0, 0];
    match reader.read_exact(&mut buf).map(|_| {
        let data: u32 = unsafe { mem::transmute(buf) };
        data
    }) {
        Err(reason) => panic!("failed to read u32: {}", Error::description(&reason)),
        Ok(byte)    => byte,
    }
}

fn read_mnist(labels_path: &Path, labels_checknum: u32, data_path: &Path, data_checknum: u32, batch_size: Option<usize>, sample_size: Option<usize>) -> Vec<(Tensor<f32>, Tensor<f32>)> {
    let labels_file = match File::open(&labels_path) {
        Err(reason) => panic!("failed to open {}: {}", labels_path.display(), Error::description(&reason)),
        Ok(file)    => file,
    };
    let data_file = match File::open(&data_path) {
        Err(reason) => panic!("failed to open {}: {}", data_path.display(), Error::description(&reason)),
        Ok(file)    => file,
    };
    let ref mut labels_reader = BufReader::new(labels_file);
    let ref mut data_reader = BufReader::new(data_file);

    let labels_magic = u32::from_be(read_u32(labels_reader));
    let data_magic = u32::from_be(read_u32(data_reader));

    assert_eq!(labels_magic, labels_checknum);
    assert_eq!(data_magic, data_checknum);

    let labels_count = u32::from_be(read_u32(labels_reader)) as usize;
    let data_count = u32::from_be(read_u32(data_reader)) as usize;
    let sample_count = cmp::min(labels_count, data_count);
    let sample_count = cmp::min(sample_count, sample_size.unwrap_or(sample_count));

    let batch_count = cmp::min(sample_count, batch_size.unwrap_or(sample_count));

    let rows = u32::from_be(read_u32(data_reader)) as usize;
    let columns = u32::from_be(read_u32(data_reader)) as usize;

    let mut sample_vec: Vec<(Tensor<f32>, Tensor<f32>)> = Vec::with_capacity(sample_count/batch_count);
    for _ in 0..sample_count/batch_count {
        let mut sample_data = vec![0u8; batch_count*rows*columns];
        match data_reader.read_exact(sample_data.as_mut()) {
            Err(reason) => panic!("failed to read data byte array: {}", Error::description(&reason)),
            Ok(_)       => (),
        }
        let mut sample_labels_byte = vec![0u8; batch_count];
        match labels_reader.read_exact(sample_labels_byte.as_mut()) {
            Err(reason) => panic!("failed to read labels byte array: {}", Error::description(&reason)),
            Ok(_)       => (),
        }

        let sample_data = sample_data.iter().map(|&x| (x as f32 - 128.0) / 74.0).collect();
        let mut sample_labels: Vec<f32> = Vec::with_capacity(batch_count * 10);
        for i in sample_labels_byte {
            match i {
                0 => sample_labels.extend([1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0].iter().cloned()),
                1 => sample_labels.extend([0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0].iter().cloned()),
                2 => sample_labels.extend([0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0].iter().cloned()),
                3 => sample_labels.extend([0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0].iter().cloned()),
                4 => sample_labels.extend([0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0].iter().cloned()),
                5 => sample_labels.extend([0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0].iter().cloned()),
                6 => sample_labels.extend([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0].iter().cloned()),
                7 => sample_labels.extend([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0].iter().cloned()),
                8 => sample_labels.extend([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0].iter().cloned()),
                9 => sample_labels.extend([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0].iter().cloned()),
                _ => (),
            }
        }

        sample_vec.push((Tensor::from_vec(Vec2(batch_count, rows*columns), sample_data), Tensor::from_vec(Vec2(batch_count, 10), sample_labels)));
    }

    sample_vec
}

#[test]
#[ignore]
fn mnist(){
    ///////////////
    // Variables //
    ///////////////

    let input_x = Arc::new(Variable::new("input_x".to_string(), Vec2(0, 28 * 28)));
    let target_y = Arc::new(Variable::new("target_y".to_string(), Vec2(0, 10)));

    // Initialize
    let mut variable_context = Context::<f32>::with_capacity(2);
    Variable::init_f32(vec![input_x.clone(), target_y.clone()], &mut variable_context);

    ///////////
    // Graph //
    ///////////

    let layers: usize = 4;
    let mut states = Vec::<Arc<State>>::with_capacity(2 * layers);
    let mut graph_head: Arc<Graph<f32>> = input_x.clone();

    {
        let w = Arc::new(State::new(format!("weight_w_{}", 1), Vec2(28 * 28, 64)));
        let b = Arc::new(State::new(format!("weight_b_{}", 1), Vec2(1, 64)));

        let dot = Arc::new(k::op::dot::<f32>(format!("layer_{}_dot", 1), graph_head.clone(), w.clone()));
        let add = Arc::new(k::op::add::<f32>(format!("layer_{}_add", 1), dot, b.clone()));

        let relu = Arc::new(k::op::relu_f32(format!("layer_{}_relu", 1), add));

        graph_head = relu;

        states.push(w);
        states.push(b);
    }

    {
        let w = Arc::new(State::new(format!("weight_w_{}", 3), Vec2(64, 32)));
        let b = Arc::new(State::new(format!("weight_b_{}", 3), Vec2(1, 32)));

        let dot = Arc::new(k::op::dot::<f32>(format!("layer_{}_dot", 3), graph_head.clone(), w.clone()));
        let add = Arc::new(k::op::add::<f32>(format!("layer_{}_add", 3), dot, b.clone()));

        let relu = Arc::new(k::op::relu_f32(format!("layer_{}_relu", 3), add));

        graph_head = relu;

        states.push(w);
        states.push(b);
    }

    {
        let w = Arc::new(State::new(format!("weight_w_{}", 4), Vec2(32, 10)));
        let b = Arc::new(State::new(format!("weight_b_{}", 4), Vec2(1, 10)));

        let dot = Arc::new(k::op::dot::<f32>(format!("layer_{}_dot", 4), graph_head.clone(), w.clone()));
        let add = Arc::new(k::op::add::<f32>(format!("layer_{}_add", 4), dot, b.clone()));

        let relu = Arc::new(k::op::relu_f32(format!("layer_{}_relu", 4), add));

        graph_head = relu;

        states.push(w);
        states.push(b);
    }

    let softmax = Arc::new(k::op::softmax_f32(format!("layer_{}_softmax", 3), graph_head.clone()));
    let xentropy = Arc::new(k::cost::softmax_cross_entropy_f32(format!("layer_{}_xentropy", 3), softmax.clone(), target_y.clone()));

    // initialize states
    let mut state_context = Context::<f32>::with_capacity(2 * layers);
    State::init_f32(states, &mut state_context);

    //////////////
    // Training //
    //////////////
    {
        // Parameters
        let learning_rate = -0.001;
        let iterations = 2048;
        let batch_size = 16;
        let sample_size = iterations * batch_size;
        let print_rate = 64;

        let train_labels_path = Path::new("data/train-labels-idx1-ubyte");
        let train_data_path = Path::new("data/train-images-idx3-ubyte");
        let training_vec = read_mnist(&train_labels_path, 2049, &train_data_path, 2051, Some(batch_size), Some(sample_size));

        let mut history = Context::<f32>::with_capacity(5 * layers + 4);

        let training_vec_length = training_vec.len();
        for i in 0..iterations {
            let (ref a, ref b) = training_vec[i % training_vec_length];
            variable_context.set(input_x.get_id(), a.clone());
            variable_context.set(target_y.get_id(), b.clone());

            k::train(xentropy.clone(), &mut state_context, &variable_context, &mut history, learning_rate);

            // test print
            if i % print_rate == 0 {
                println!("iteration: {} | cross entropy cost: {}", i,
                k::execute(xentropy.clone(), &state_context, &variable_context).get(Vec2(0, 0)));
            }
        }

        println!("\ntraining complete\nalpha: {}, iterations: {}, batch_size: {}", learning_rate, iterations, batch_size);
    }

    ////////////////
    // Final test //
    ////////////////
    {
        println!("\nFinal Test:\n");
        // Parameters
        let batch_size = None;
        let sample_size = None;

        let test_labels_path = Path::new("data/t10k-labels-idx1-ubyte");
        let test_data_path = Path::new("data/t10k-images-idx3-ubyte");
        let test_vec = read_mnist(&test_labels_path, 2049, &test_data_path, 2051, batch_size, sample_size);

        let (ref a, ref b) = test_vec[0];
        variable_context.set(input_x.get_id(), a.clone());
        variable_context.set(target_y.get_id(), b.clone());

        let result_tensor = k::op::softmax_round_f32(k::execute(softmax.clone(), &state_context, &variable_context)) + (k::op::softmax_round_f32(b.clone()) * -1);
        let total = result_tensor.buffer().len() as f64;
        let mut score = 0.0;
        for &i in result_tensor.buffer().iter() {
            if i == 0 {
                score += 1.0;
            }
        }

        println!("{}", score/total);
    }
}
<commit_msg>simplify mnist<commit_after>extern crate ktensor as k;
use k::{Arc, Vec2, Tensor, Context, Graph, State, Variable};
use std::path::{Path};
use std::fs::{File};
use std::io::{BufReader, Read};
use std::error::{Error};
use std::cmp;

fn read_u8(reader: &mut Read) -> u8 {
    use std::mem;

    let mut buf: [u8; 1] = [0];
    match reader.read_exact(&mut buf).map(|_| {
        let data: u8 = unsafe { mem::transmute(buf) };
        data
    }) {
        Err(reason) => panic!("failed to read u8: {}", Error::description(&reason)),
        Ok(byte)    => byte,
    }
}

fn read_u32(reader: &mut Read) -> u32 {
    use std::mem;

    let mut buf: [u8; 4] = [0, 0, 0, 0];
    match reader.read_exact(&mut buf).map(|_| {
        let data: u32 = unsafe { mem::transmute(buf) };
        data
    }) {
        Err(reason) => panic!("failed to read u32: {}", Error::description(&reason)),
        Ok(byte)    => byte,
    }
}

fn read_mnist(labels_path: &Path, labels_checknum: u32, data_path: &Path, data_checknum: u32, batch_size: Option<usize>, sample_size: Option<usize>) -> Vec<(Tensor<f32>, Tensor<f32>)> {
    let labels_file = match File::open(&labels_path) {
        Err(reason) => panic!("failed to open {}: {}", labels_path.display(), Error::description(&reason)),
        Ok(file)    => file,
    };
    let data_file = match File::open(&data_path) {
        Err(reason) => panic!("failed to open {}: {}", data_path.display(), Error::description(&reason)),
        Ok(file)    => file,
    };
    let ref mut labels_reader = BufReader::new(labels_file);
    let ref mut data_reader = BufReader::new(data_file);

    let labels_magic = u32::from_be(read_u32(labels_reader));
    let data_magic = u32::from_be(read_u32(data_reader));

    assert_eq!(labels_magic, labels_checknum);
    assert_eq!(data_magic, data_checknum);

    let labels_count = u32::from_be(read_u32(labels_reader)) as usize;
    let data_count = u32::from_be(read_u32(data_reader)) as usize;
    let sample_count = cmp::min(labels_count, data_count);
    let sample_count = cmp::min(sample_count, sample_size.unwrap_or(sample_count));

    let batch_count = cmp::min(sample_count, batch_size.unwrap_or(sample_count));

    let rows = u32::from_be(read_u32(data_reader)) as usize;
    let columns = u32::from_be(read_u32(data_reader)) as usize;

    let mut sample_vec: Vec<(Tensor<f32>, Tensor<f32>)> = Vec::with_capacity(sample_count/batch_count);
    for _ in 0..sample_count/batch_count {
        let mut sample_data = vec![0u8; batch_count*rows*columns];
        match data_reader.read_exact(sample_data.as_mut()) {
            Err(reason) => panic!("failed to read data byte array: {}", Error::description(&reason)),
            Ok(_)       => (),
        }
        let mut sample_labels_byte = vec![0u8; batch_count];
        match labels_reader.read_exact(sample_labels_byte.as_mut()) {
            Err(reason) => panic!("failed to read labels byte array: {}", Error::description(&reason)),
            Ok(_)       => (),
        }

        let sample_data = sample_data.iter().map(|&x| (x as f32 - 128.0) / 74.0).collect();
        let mut sample_labels: Vec<f32> = Vec::with_capacity(batch_count * 10);
        for i in sample_labels_byte {
            match i {
                0 => sample_labels.extend([1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0].iter().cloned()),
                1 => sample_labels.extend([0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0].iter().cloned()),
                2 => sample_labels.extend([0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0].iter().cloned()),
                3 => sample_labels.extend([0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0].iter().cloned()),
                4 => sample_labels.extend([0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0].iter().cloned()),
                5 => sample_labels.extend([0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0].iter().cloned()),
                6 => sample_labels.extend([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0].iter().cloned()),
                7 => sample_labels.extend([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0].iter().cloned()),
                8 => sample_labels.extend([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0].iter().cloned()),
                9 => sample_labels.extend([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0].iter().cloned()),
                _ => (),
            }
        }

        sample_vec.push((Tensor::from_vec(Vec2(batch_count, rows*columns), sample_data), Tensor::from_vec(Vec2(batch_count, 10), sample_labels)));
    }

    sample_vec
}

#[test]
#[ignore]
fn mnist(){
    ///////////////
    // Variables //
    ///////////////

    let input_x = Arc::new(Variable::new("input_x".to_string(), Vec2(0, 28 * 28)));
    let target_y = Arc::new(Variable::new("target_y".to_string(), Vec2(0, 10)));

    // Initialize
    let mut variable_context = Context::<f32>::with_capacity(2);
    Variable::init_f32(vec![input_x.clone(), target_y.clone()], &mut variable_context);

    ///////////
    // Graph //
    ///////////

    let layers: usize = 2;
    let mut states = Vec::<Arc<State>>::with_capacity(2 * layers);
    let mut graph_head: Arc<Graph<f32>> = input_x.clone();

    {
        let w = Arc::new(State::new(format!("weight_w_{}", 1), Vec2(28 * 28, 64)));
        let b = Arc::new(State::new(format!("weight_b_{}", 1), Vec2(1, 64)));

        let dot = Arc::new(k::op::dot::<f32>(format!("layer_{}_dot", 1), graph_head.clone(), w.clone()));
        let add = Arc::new(k::op::add::<f32>(format!("layer_{}_add", 1), dot, b.clone()));

        let relu = Arc::new(k::op::relu_f32(format!("layer_{}_relu", 1), add));

        graph_head = relu;

        states.push(w);
        states.push(b);
    }

    {
        let w = Arc::new(State::new(format!("weight_w_{}", 2), Vec2(64, 10)));
        let b = Arc::new(State::new(format!("weight_b_{}", 2), Vec2(1, 10)));

        let dot = Arc::new(k::op::dot::<f32>(format!("layer_{}_dot", 2), graph_head.clone(), w.clone()));
        let add = Arc::new(k::op::add::<f32>(format!("layer_{}_add", 2), dot, b.clone()));

        let relu = Arc::new(k::op::relu_f32(format!("layer_{}_relu", 2), add));

        graph_head = relu;

        states.push(w);
        states.push(b);
    }

    let softmax = Arc::new(k::op::softmax_f32(format!("layer_{}_softmax", 3), graph_head.clone()));
    let xentropy = Arc::new(k::cost::softmax_cross_entropy_f32(format!("layer_{}_xentropy", 3), softmax.clone(), target_y.clone()));

    // initialize states
    let mut state_context = Context::<f32>::with_capacity(2 * layers);
    State::init_f32(states, &mut state_context);

    //////////////
    // Training //
    //////////////
    {
        // Parameters
        let learning_rate = -0.001;
        let iterations = 4096;
        let batch_size = 16;
        let sample_size = iterations * batch_size;
        let print_rate = 64;

        let train_labels_path = Path::new("data/train-labels-idx1-ubyte");
        let train_data_path = Path::new("data/train-images-idx3-ubyte");
        let training_vec = read_mnist(&train_labels_path, 2049, &train_data_path, 2051, Some(batch_size), Some(sample_size));

        let mut history = Context::<f32>::with_capacity(5 * layers + 4);

        let training_vec_length = training_vec.len();
        for i in 0..iterations {
            let (ref a, ref b) = training_vec[i % training_vec_length];
            variable_context.set(input_x.get_id(), a.clone());
            variable_context.set(target_y.get_id(), b.clone());

            k::train(xentropy.clone(), &mut state_context, &variable_context, &mut history, learning_rate);

            // test print
            if i % print_rate == 0 {
                println!("iteration: {} | cross entropy cost: {}", i,
                k::execute(xentropy.clone(), &state_context, &variable_context).get(Vec2(0, 0)));
            }
        }

        println!("\ntraining complete\nalpha: {}, iterations: {}, batch_size: {}", learning_rate, iterations, batch_size);
    }

    ////////////////
    // Final test //
    ////////////////
    {
        println!("\nFinal Test:\n");
        // Parameters
        let batch_size = None;
        let sample_size = None;

        let test_labels_path = Path::new("data/t10k-labels-idx1-ubyte");
        let test_data_path = Path::new("data/t10k-images-idx3-ubyte");
        let test_vec = read_mnist(&test_labels_path, 2049, &test_data_path, 2051, batch_size, sample_size);

        let (ref a, ref b) = test_vec[0];
        variable_context.set(input_x.get_id(), a.clone());
        variable_context.set(target_y.get_id(), b.clone());

        let result_tensor = k::op::softmax_round_f32(k::execute(softmax.clone(), &state_context, &variable_context)) + (k::op::softmax_round_f32(b.clone()) * -1);
        let total = result_tensor.buffer().len() as f64;
        let mut score = 0.0;
        for &i in result_tensor.buffer().iter() {
            if i == 0 {
                score += 1.0;
            }
        }

        println!("{}", score/total);
    }
}
<|endoftext|>