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

use std::time::*;
use std::sync::Arc;

use fail;
use futures::Future;
use kvproto::raft_serverpb::{PeerState, RegionLocalState};
use tikv::pd::PdClient;
use tikv::raftstore::store::keys;
use tikv::storage::CF_RAFT;
use tikv::raftstore::store::Peekable;

use raftstore::node::new_node_cluster;
use raftstore::transport_simulate::*;
use raftstore::util;
use raftstore::util::*;

/// Test if merge is rollback as expected.
#[test]
fn test_node_merge_rollback() {
    let _guard = ::setup();
    let mut cluster = new_node_cluster(0, 3);
    configure_for_merge(&mut cluster);
    let pd_client = Arc::clone(&cluster.pd_client);
    pd_client.disable_default_operator();

    cluster.run_conf_change();

    let region = pd_client.get_region(b"k1").unwrap();
    cluster.must_split(&region, b"k2");
    let left = pd_client.get_region(b"k1").unwrap();
    let right = pd_client.get_region(b"k2").unwrap();

    pd_client.must_add_peer(left.get_id(), new_peer(2, 2));
    pd_client.must_add_peer(right.get_id(), new_peer(2, 4));

    cluster.must_put(b"k1", b"v1");
    cluster.must_put(b"k3", b"v3");

    let region = pd_client.get_region(b"k1").unwrap();
    let target_region = pd_client.get_region(b"k3").unwrap();

    let schedule_merge_fp = "on_schedule_merge";
    fail::cfg(schedule_merge_fp, "return()").unwrap();

    // The call is finished when prepare_merge is applied.
    cluster.try_merge(region.get_id(), target_region.get_id());

    // Add a peer to trigger rollback.
    pd_client.must_add_peer(right.get_id(), new_peer(3, 5));
    cluster.must_put(b"k4", b"v4");
    util::must_get_equal(&cluster.get_engine(3), b"k4", b"v4");

    let mut region = pd_client.get_region(b"k1").unwrap();
    // After split and prepare_merge, version becomes 1 + 2 = 3;
    assert_eq!(region.get_region_epoch().get_version(), 3);
    // After ConfChange and prepare_merge, conf version becomes 1 + 2 = 3;
    assert_eq!(region.get_region_epoch().get_conf_ver(), 3);
    fail::remove(schedule_merge_fp);
    // Wait till rollback.
    cluster.must_put(b"k11", b"v11");

    // After rollback, version becomes 3 + 1 = 4;
    region.mut_region_epoch().set_version(4);
    for i in 1..3 {
        let state_key = keys::region_state_key(region.get_id());
        let state: RegionLocalState = cluster
            .get_engine(i)
            .get_msg_cf(CF_RAFT, &state_key)
            .unwrap()
            .unwrap();
        assert_eq!(state.get_state(), PeerState::Normal);
        assert_eq!(*state.get_region(), region);
    }

    pd_client.must_remove_peer(right.get_id(), new_peer(3, 5));
    fail::cfg(schedule_merge_fp, "return()").unwrap();

    let target_region = pd_client.get_region(b"k1").unwrap();
    cluster.try_merge(region.get_id(), target_region.get_id());
    let mut region = pd_client.get_region(b"k1").unwrap();

    // Split to trigger rollback.
    cluster.must_split(&right, b"k3");
    fail::remove(schedule_merge_fp);
    // Wait till rollback.
    cluster.must_put(b"k12", b"v12");

    // After premerge and rollback, version becomes 4 + 2 = 6;
    region.mut_region_epoch().set_version(4);
    for i in 1..3 {
        let state_key = keys::region_state_key(region.get_id());
        let state: RegionLocalState = cluster
            .get_engine(i)
            .get_msg_cf(CF_RAFT, &state_key)
            .unwrap()
            .unwrap();
        assert_eq!(state.get_state(), PeerState::Normal);
        assert_eq!(*state.get_region(), region);
    }
}

/// Test if merge is still working when restart a cluster during merge.
#[test]
fn test_node_merge_restart() {
    let _guard = ::setup();
    // ::util::init_log();
    let mut cluster = new_node_cluster(0, 3);
    configure_for_merge(&mut cluster);
    cluster.run();

    let pd_client = Arc::clone(&cluster.pd_client);
    let region = pd_client.get_region(b"k1").unwrap();
    cluster.must_split(&region, b"k2");
    let left = pd_client.get_region(b"k1").unwrap();
    let right = pd_client.get_region(b"k2").unwrap();

    cluster.must_put(b"k1", b"v1");
    cluster.must_put(b"k3", b"v3");

    let schedule_merge_fp = "on_schedule_merge";
    fail::cfg(schedule_merge_fp, "return()").unwrap();

    cluster.try_merge(left.get_id(), right.get_id());
    let leader = cluster.leader_of_region(left.get_id()).unwrap();

    cluster.shutdown();
    let engine = cluster.get_engine(leader.get_store_id());
    let state_key = keys::region_state_key(left.get_id());
    let state: RegionLocalState = engine.get_msg_cf(CF_RAFT, &state_key).unwrap().unwrap();
    assert_eq!(state.get_state(), PeerState::Merging, "{:?}", state);
    let state_key = keys::region_state_key(right.get_id());
    let state: RegionLocalState = engine.get_msg_cf(CF_RAFT, &state_key).unwrap().unwrap();
    assert_eq!(state.get_state(), PeerState::Normal, "{:?}", state);
    fail::remove(schedule_merge_fp);
    cluster.start();

    // Wait till merge is finished.
    let timer = Instant::now();
    loop {
        if pd_client
            .get_region_by_id(left.get_id())
            .wait()
            .unwrap()
            .is_none()
        {
            break;
        }

        if timer.elapsed() > Duration::from_secs(5) {
            panic!("region still not merged after 5 secs");
        }
        sleep_ms(10);
    }

    cluster.must_put(b"k4", b"v4");

    for i in 1..4 {
        util::must_get_equal(&cluster.get_engine(i), b"k4", b"v4");
        let state_key = keys::region_state_key(left.get_id());
        let state: RegionLocalState = cluster
            .get_engine(i)
            .get_msg_cf(CF_RAFT, &state_key)
            .unwrap()
            .unwrap();
        assert_eq!(state.get_state(), PeerState::Tombstone, "{:?}", state);
        let state_key = keys::region_state_key(right.get_id());
        let state: RegionLocalState = cluster
            .get_engine(i)
            .get_msg_cf(CF_RAFT, &state_key)
            .unwrap()
            .unwrap();
        assert_eq!(state.get_state(), PeerState::Normal, "{:?}", state);
        assert!(state.get_region().get_start_key().is_empty());
        assert!(state.get_region().get_end_key().is_empty());
    }

    // Now test if cluster works fine when it crash after merge is applied
    // but before notifying raftstore thread.
    let region = pd_client.get_region(b"k1").unwrap();
    let peer_on_store1 = find_peer(&region, 1).unwrap().to_owned();
    cluster.must_transfer_leader(region.get_id(), peer_on_store1);
    cluster.must_split(&region, b"k2");
    let left = pd_client.get_region(b"k1").unwrap();
    let right = pd_client.get_region(b"k2").unwrap();
    cluster.must_put(b"k11", b"v11");
    must_get_equal(&cluster.get_engine(3), b"k11", b"v11");
    let skip_destroy_fp = "raft_store_skip_destroy_peer";
    fail::cfg(skip_destroy_fp, "return()").unwrap();
    cluster.add_send_filter(IsolationFilterFactory::new(3));
    pd_client.must_merge(left.get_id(), right.get_id());
    let peer = find_peer(&right, 3).unwrap().to_owned();
    pd_client.must_remove_peer(right.get_id(), peer);
    cluster.shutdown();
    fail::remove(skip_destroy_fp);
    cluster.clear_send_filters();
    cluster.start();
    must_get_none(&cluster.get_engine(3), b"k1");
    must_get_none(&cluster.get_engine(3), b"k3");
}

/// Test if merging state will be removed after accepting a snapshot.
#[test]
fn test_node_merge_recover_snapshot() {
    let _guard = ::setup();
    let mut cluster = new_node_cluster(0, 3);
    configure_for_merge(&mut cluster);
    cluster.cfg.raft_store.raft_log_gc_threshold = 12;
    cluster.cfg.raft_store.raft_log_gc_count_limit = 12;
    let pd_client = Arc::clone(&cluster.pd_client);
    pd_client.disable_default_operator();

    cluster.run();

    let region = pd_client.get_region(b"k1").unwrap();
    cluster.must_split(&region, b"k2");
    let left = pd_client.get_region(b"k1").unwrap();

    cluster.must_put(b"k1", b"v1");
    cluster.must_put(b"k3", b"v3");

    let region = pd_client.get_region(b"k3").unwrap();
    let target_region = pd_client.get_region(b"k1").unwrap();

    let schedule_merge_fp = "on_schedule_merge";
    fail::cfg(schedule_merge_fp, "return()").unwrap();

    cluster.try_merge(region.get_id(), target_region.get_id());

    // Remove a peer to trigger rollback.
    pd_client.must_remove_peer(left.get_id(), left.get_peers()[0].to_owned());
    util::must_get_none(&cluster.get_engine(3), b"k4");

    let step_store_3_region_1 = "step_message_3_1";
    fail::cfg(step_store_3_region_1, "return()").unwrap();
    fail::remove(schedule_merge_fp);

    for i in 0..100 {
        cluster.must_put(format!("k4{}", i).as_bytes(), b"v4");
    }
    fail::remove(step_store_3_region_1);
    util::must_get_equal(&cluster.get_engine(3), b"k40", b"v4");
    cluster.must_transfer_leader(1, new_peer(3, 3));
    cluster.must_put(b"k40", b"v5");
}
<commit_msg>tests: make merge rollback tests more stable (#2903)<commit_after>// Copyright 2017 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

use std::time::*;
use std::sync::Arc;

use fail;
use futures::Future;
use kvproto::raft_serverpb::{PeerState, RegionLocalState};
use tikv::pd::PdClient;
use tikv::raftstore::store::keys;
use tikv::storage::CF_RAFT;
use tikv::raftstore::store::Peekable;

use raftstore::node::new_node_cluster;
use raftstore::transport_simulate::*;
use raftstore::util;
use raftstore::util::*;

/// Test if merge is rollback as expected.
#[test]
fn test_node_merge_rollback() {
    let _guard = ::setup();
    let mut cluster = new_node_cluster(0, 3);
    configure_for_merge(&mut cluster);
    let pd_client = Arc::clone(&cluster.pd_client);
    pd_client.disable_default_operator();

    cluster.run_conf_change();

    let region = pd_client.get_region(b"k1").unwrap();
    cluster.must_split(&region, b"k2");
    let left = pd_client.get_region(b"k1").unwrap();
    let right = pd_client.get_region(b"k2").unwrap();

    pd_client.must_add_peer(left.get_id(), new_peer(2, 2));
    pd_client.must_add_peer(right.get_id(), new_peer(2, 4));

    cluster.must_put(b"k1", b"v1");
    cluster.must_put(b"k3", b"v3");

    let region = pd_client.get_region(b"k1").unwrap();
    let target_region = pd_client.get_region(b"k3").unwrap();

    let schedule_merge_fp = "on_schedule_merge";
    fail::cfg(schedule_merge_fp, "return()").unwrap();

    // The call is finished when prepare_merge is applied.
    cluster.try_merge(region.get_id(), target_region.get_id());

    // Add a peer to trigger rollback.
    pd_client.must_add_peer(right.get_id(), new_peer(3, 5));
    cluster.must_put(b"k4", b"v4");
    util::must_get_equal(&cluster.get_engine(3), b"k4", b"v4");

    let mut region = pd_client.get_region(b"k1").unwrap();
    // After split and prepare_merge, version becomes 1 + 2 = 3;
    assert_eq!(region.get_region_epoch().get_version(), 3);
    // After ConfChange and prepare_merge, conf version becomes 1 + 2 = 3;
    assert_eq!(region.get_region_epoch().get_conf_ver(), 3);
    fail::remove(schedule_merge_fp);
    // Wait till rollback.
    cluster.must_put(b"k11", b"v11");

    // After rollback, version becomes 3 + 1 = 4;
    region.mut_region_epoch().set_version(4);
    for i in 1..3 {
        must_get_equal(&cluster.get_engine(i), b"k11", b"v11");
        let state_key = keys::region_state_key(region.get_id());
        let state: RegionLocalState = cluster
            .get_engine(i)
            .get_msg_cf(CF_RAFT, &state_key)
            .unwrap()
            .unwrap();
        assert_eq!(state.get_state(), PeerState::Normal);
        assert_eq!(*state.get_region(), region);
    }

    pd_client.must_remove_peer(right.get_id(), new_peer(3, 5));
    fail::cfg(schedule_merge_fp, "return()").unwrap();

    let target_region = pd_client.get_region(b"k1").unwrap();
    cluster.try_merge(region.get_id(), target_region.get_id());
    let mut region = pd_client.get_region(b"k1").unwrap();

    // Split to trigger rollback.
    cluster.must_split(&right, b"k3");
    fail::remove(schedule_merge_fp);
    // Wait till rollback.
    cluster.must_put(b"k12", b"v12");

    // After premerge and rollback, version becomes 4 + 2 = 6;
    region.mut_region_epoch().set_version(4);
    for i in 1..3 {
        must_get_equal(&cluster.get_engine(i), b"k12", b"v12");
        let state_key = keys::region_state_key(region.get_id());
        let state: RegionLocalState = cluster
            .get_engine(i)
            .get_msg_cf(CF_RAFT, &state_key)
            .unwrap()
            .unwrap();
        assert_eq!(state.get_state(), PeerState::Normal);
        assert_eq!(*state.get_region(), region);
    }
}

/// Test if merge is still working when restart a cluster during merge.
#[test]
fn test_node_merge_restart() {
    let _guard = ::setup();
    // ::util::init_log();
    let mut cluster = new_node_cluster(0, 3);
    configure_for_merge(&mut cluster);
    cluster.run();

    let pd_client = Arc::clone(&cluster.pd_client);
    let region = pd_client.get_region(b"k1").unwrap();
    cluster.must_split(&region, b"k2");
    let left = pd_client.get_region(b"k1").unwrap();
    let right = pd_client.get_region(b"k2").unwrap();

    cluster.must_put(b"k1", b"v1");
    cluster.must_put(b"k3", b"v3");

    let schedule_merge_fp = "on_schedule_merge";
    fail::cfg(schedule_merge_fp, "return()").unwrap();

    cluster.try_merge(left.get_id(), right.get_id());
    let leader = cluster.leader_of_region(left.get_id()).unwrap();

    cluster.shutdown();
    let engine = cluster.get_engine(leader.get_store_id());
    let state_key = keys::region_state_key(left.get_id());
    let state: RegionLocalState = engine.get_msg_cf(CF_RAFT, &state_key).unwrap().unwrap();
    assert_eq!(state.get_state(), PeerState::Merging, "{:?}", state);
    let state_key = keys::region_state_key(right.get_id());
    let state: RegionLocalState = engine.get_msg_cf(CF_RAFT, &state_key).unwrap().unwrap();
    assert_eq!(state.get_state(), PeerState::Normal, "{:?}", state);
    fail::remove(schedule_merge_fp);
    cluster.start();

    // Wait till merge is finished.
    let timer = Instant::now();
    loop {
        if pd_client
            .get_region_by_id(left.get_id())
            .wait()
            .unwrap()
            .is_none()
        {
            break;
        }

        if timer.elapsed() > Duration::from_secs(5) {
            panic!("region still not merged after 5 secs");
        }
        sleep_ms(10);
    }

    cluster.must_put(b"k4", b"v4");

    for i in 1..4 {
        util::must_get_equal(&cluster.get_engine(i), b"k4", b"v4");
        let state_key = keys::region_state_key(left.get_id());
        let state: RegionLocalState = cluster
            .get_engine(i)
            .get_msg_cf(CF_RAFT, &state_key)
            .unwrap()
            .unwrap();
        assert_eq!(state.get_state(), PeerState::Tombstone, "{:?}", state);
        let state_key = keys::region_state_key(right.get_id());
        let state: RegionLocalState = cluster
            .get_engine(i)
            .get_msg_cf(CF_RAFT, &state_key)
            .unwrap()
            .unwrap();
        assert_eq!(state.get_state(), PeerState::Normal, "{:?}", state);
        assert!(state.get_region().get_start_key().is_empty());
        assert!(state.get_region().get_end_key().is_empty());
    }

    // Now test if cluster works fine when it crash after merge is applied
    // but before notifying raftstore thread.
    let region = pd_client.get_region(b"k1").unwrap();
    let peer_on_store1 = find_peer(&region, 1).unwrap().to_owned();
    cluster.must_transfer_leader(region.get_id(), peer_on_store1);
    cluster.must_split(&region, b"k2");
    let left = pd_client.get_region(b"k1").unwrap();
    let right = pd_client.get_region(b"k2").unwrap();
    cluster.must_put(b"k11", b"v11");
    must_get_equal(&cluster.get_engine(3), b"k11", b"v11");
    let skip_destroy_fp = "raft_store_skip_destroy_peer";
    fail::cfg(skip_destroy_fp, "return()").unwrap();
    cluster.add_send_filter(IsolationFilterFactory::new(3));
    pd_client.must_merge(left.get_id(), right.get_id());
    let peer = find_peer(&right, 3).unwrap().to_owned();
    pd_client.must_remove_peer(right.get_id(), peer);
    cluster.shutdown();
    fail::remove(skip_destroy_fp);
    cluster.clear_send_filters();
    cluster.start();
    must_get_none(&cluster.get_engine(3), b"k1");
    must_get_none(&cluster.get_engine(3), b"k3");
}

/// Test if merging state will be removed after accepting a snapshot.
#[test]
fn test_node_merge_recover_snapshot() {
    let _guard = ::setup();
    let mut cluster = new_node_cluster(0, 3);
    configure_for_merge(&mut cluster);
    cluster.cfg.raft_store.raft_log_gc_threshold = 12;
    cluster.cfg.raft_store.raft_log_gc_count_limit = 12;
    let pd_client = Arc::clone(&cluster.pd_client);
    pd_client.disable_default_operator();

    cluster.run();

    let region = pd_client.get_region(b"k1").unwrap();
    cluster.must_split(&region, b"k2");
    let left = pd_client.get_region(b"k1").unwrap();

    cluster.must_put(b"k1", b"v1");
    cluster.must_put(b"k3", b"v3");

    let region = pd_client.get_region(b"k3").unwrap();
    let target_region = pd_client.get_region(b"k1").unwrap();

    let schedule_merge_fp = "on_schedule_merge";
    fail::cfg(schedule_merge_fp, "return()").unwrap();

    cluster.try_merge(region.get_id(), target_region.get_id());

    // Remove a peer to trigger rollback.
    pd_client.must_remove_peer(left.get_id(), left.get_peers()[0].to_owned());
    util::must_get_none(&cluster.get_engine(3), b"k4");

    let step_store_3_region_1 = "step_message_3_1";
    fail::cfg(step_store_3_region_1, "return()").unwrap();
    fail::remove(schedule_merge_fp);

    for i in 0..100 {
        cluster.must_put(format!("k4{}", i).as_bytes(), b"v4");
    }
    fail::remove(step_store_3_region_1);
    util::must_get_equal(&cluster.get_engine(3), b"k40", b"v4");
    cluster.must_transfer_leader(1, new_peer(3, 3));
    cluster.must_put(b"k40", b"v5");
}
<|endoftext|>


<commit_before>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

#![forbid(unsafe_code)]

use crate::{
    cluster::Cluster,
    experiments::{Context, Experiment, ExperimentParam},
    instance,
    instance::Instance,
    tx_emitter::EmitJobRequest,
};
use async_trait::async_trait;
use diem_logger::prelude::*;
use futures::future::try_join_all;
use std::{
    collections::HashSet,
    env, fmt,
    iter::once,
    time::{Duration, Instant},
};
use structopt::StructOpt;
use tokio::time;

/// Reboot `updated_instance` with newer image tag
pub async fn update_batch_instance(
    context: &mut Context<'_>,
    updated_instance: &[Instance],
    updated_lsr: &[Instance],
    updated_tag: String,
) -> anyhow::Result<()> {
    info!("Stop Existing instances.");
    let futures: Vec<_> = updated_instance.iter().map(Instance::stop).collect();
    try_join_all(futures).await?;

    if !updated_lsr.is_empty() {
        info!("Stop associated lsr instances.");
        let futures: Vec<_> = updated_lsr.iter().map(Instance::stop).collect();
        try_join_all(futures).await?;
        info!("Reinstantiate a set of new lsr.");
        let futures: Vec<_> = updated_lsr
            .iter()
            .map(|instance| {
                let mut newer_config = instance.instance_config().clone();
                newer_config.replace_tag(updated_tag.clone()).unwrap();
                context.cluster_swarm.spawn_new_instance(newer_config)
            })
            .collect();
        try_join_all(futures).await?;
        info!("Wait for the instance to sync up with peers");
        time::sleep(Duration::from_secs(20)).await;
    }

    info!("Reinstantiate a set of new nodes.");
    let futures: Vec<_> = updated_instance
        .iter()
        .map(|instance| {
            let mut newer_config = instance.instance_config().clone();
            newer_config.replace_tag(updated_tag.clone()).unwrap();
            context.cluster_swarm.spawn_new_instance(newer_config)
        })
        .collect();
    let instances = try_join_all(futures).await?;

    info!("Wait for the instances to recover.");
    let deadline = Instant::now() + Duration::from_secs(5 * 60);
    let futures: Vec<_> = instances
        .iter()
        .map(|instance| instance.wait_json_rpc(deadline))
        .collect();
    try_join_all(futures).await?;

    // Add a timeout to have wait for validators back to healthy mode.
    // TODO: Replace this with a blocking health check.
    info!("Wait for the instance to sync up with peers");
    time::sleep(Duration::from_secs(20)).await;
    Ok(())
}

pub fn get_instance_list_str(batch: &[Instance]) -> String {
    let mut nodes_list = String::from("");
    for instance in batch.iter() {
        nodes_list.push_str(&instance.to_string());
        nodes_list.push_str(", ")
    }
    nodes_list
}

#[derive(StructOpt, Debug)]
pub struct CompatiblityTestParams {
    #[structopt(
        long,
        default_value = "15",
        help = "Number of nodes to update in the first batch"
    )]
    pub count: usize,
    #[structopt(long, help = "Image tag of newer validator software")]
    pub updated_image_tag: String,
}

pub struct CompatibilityTest {
    first_node: Instance,
    first_lsr: Vec<Instance>,
    first_batch: Vec<Instance>,
    first_batch_lsr: Vec<Instance>,
    second_batch: Vec<Instance>,
    second_batch_lsr: Vec<Instance>,
    full_nodes: Vec<Instance>,
    updated_image_tag: String,
}

impl ExperimentParam for CompatiblityTestParams {
    type E = CompatibilityTest;
    fn build(self, cluster: &Cluster) -> Self::E {
        if self.count > cluster.validator_instances().len() || self.count == 0 {
            panic!(
                "Can not reboot {} validators in cluster with {} instances",
                self.count,
                cluster.validator_instances().len()
            );
        }
        let (first_batch, second_batch) = cluster.split_n_validators_random(self.count);
        let mut first_batch = first_batch.into_validator_instances();
        let second_batch = second_batch.into_validator_instances();
        let first_node = first_batch
            .pop()
            .expect("Requires at least one validator in the first batch");
        let mut first_lsr = vec![];
        let mut first_batch_lsr = vec![];
        let mut second_batch_lsr = vec![];
        if !cluster.lsr_instances().is_empty() {
            first_batch_lsr = cluster.lsr_instances_for_validators(&first_batch);
            second_batch_lsr = cluster.lsr_instances_for_validators(&second_batch);
            first_lsr = cluster.lsr_instances_for_validators(&[first_node.clone()]);
        }

        Self::E {
            first_node,
            first_lsr,
            first_batch,
            first_batch_lsr,
            second_batch,
            second_batch_lsr,
            full_nodes: cluster.fullnode_instances().to_vec(),
            updated_image_tag: self.updated_image_tag,
        }
    }
}

#[async_trait]
impl Experiment for CompatibilityTest {
    fn affected_validators(&self) -> HashSet<String> {
        instance::instancelist_to_set(&self.first_batch)
            .union(&instance::instancelist_to_set(&self.second_batch))
            .cloned()
            .chain(once(self.first_node.peer_name().clone()))
            .collect()
    }

    async fn run(&mut self, context: &mut Context<'_>) -> anyhow::Result<()> {
        let validator_txn_job = EmitJobRequest::for_instances(
            context.cluster.validator_instances().to_vec(),
            context.global_emit_job_request,
            0,
            0,
        );
        let fullnode_txn_job = EmitJobRequest::for_instances(
            context.cluster.fullnode_instances().to_vec(),
            context.global_emit_job_request,
            0,
            0,
        );
        let job_duration = Duration::from_secs(3);
        context.report.report_text(format!(
            "Compatibility test results for {} ==> {} (PR)",
            context.current_tag, self.updated_image_tag
        ));

        // Generate some traffic
        let msg = format!(
            "1. All instances running {}, generating some traffic on network",
            context.current_tag
        );
        info!("{}", msg);
        context.report.report_text(msg);
        context
            .tx_emitter
            .emit_txn_for(job_duration, fullnode_txn_job.clone())
            .await
            .map_err(|e| anyhow::format_err!("Failed to generate traffic: {}", e))?;

        let msg = format!(
            "2. First validator {} ==> {}, to validate storage",
            context.current_tag, self.updated_image_tag
        );
        info!("{}", msg);
        info!("Upgrading validator: {}", self.first_node);
        context.report.report_text(msg);
        let first_node = vec![self.first_node.clone()];
        update_batch_instance(
            context,
            &first_node,
            &self.first_lsr,
            self.updated_image_tag.clone(),
        )
        .await?;
        context
            .tx_emitter
            .emit_txn_for(
                job_duration,
                EmitJobRequest::for_instances(first_node, context.global_emit_job_request, 0, 0),
            )
            .await
            .map_err(|e| anyhow::format_err!("Storage backwards compat broken: {}", e))?;

        let msg = format!(
            "3. First batch validators ({}) {} ==> {}, to test consensus",
            self.first_batch.len(),
            context.current_tag,
            self.updated_image_tag
        );
        info!("{}", msg);
        info!(
            "Upgrading validators: {}",
            get_instance_list_str(&self.first_batch)
        );
        context.report.report_text(msg);
        update_batch_instance(
            context,
            &self.first_batch,
            &self.first_batch_lsr,
            self.updated_image_tag.clone(),
        )
        .await?;
        context
            .tx_emitter
            .emit_txn_for(job_duration, validator_txn_job.clone())
            .await
            .map_err(|e| anyhow::format_err!("Consensus backwards compat broken: {}", e))?;

        let msg = format!(
            "4. Second batch validators ({}) {} ==> {}, to upgrade rest of the validators",
            self.second_batch.len(),
            context.current_tag,
            self.updated_image_tag
        );
        info!("{}", msg);
        info!(
            "Upgrading validators: {}",
            get_instance_list_str(&self.second_batch)
        );
        context.report.report_text(msg);
        update_batch_instance(
            context,
            &self.second_batch,
            &self.second_batch_lsr,
            self.updated_image_tag.clone(),
        )
        .await?;
        context
            .tx_emitter
            .emit_txn_for(job_duration, validator_txn_job)
            .await
            .map_err(|e| {
                anyhow::format_err!("Failed to upgrade rest of validator images: {}", e)
            })?;

        // TODO(rustielin): fullnode reboot in cluster-test breaks with current version of config-builder
        //                  skipping for now, so cluster is not restored to same version across all instances
        let disable_fn_upgrade = env::var("DISABLE_FN_UPGRADE").is_ok();
        if disable_fn_upgrade {
            let msg = format!(
                "5. Reset all nodes ==> {}, to finish the network upgrade",
                self.updated_image_tag
            );
            info!("{}", msg);
            context.report.report_text(msg);
            context.cluster_builder.current_tag = self.updated_image_tag.clone();
            context
                .cluster_builder
                .setup_cluster(context.cluster_builder_params, false)
                .await?;
        } else {
            let msg = format!(
                "5. All full nodes ({}) {} ==> {}, to finish the network upgrade",
                self.full_nodes.len(),
                context.current_tag,
                self.updated_image_tag
            );
            info!("{}", msg);
            info!(
                "Upgrading full nodes: {}",
                get_instance_list_str(&self.full_nodes)
            );
            context.report.report_text(msg);
            update_batch_instance(
                context,
                &self.full_nodes,
                &[],
                self.updated_image_tag.clone(),
            )
            .await?;
            context
                .tx_emitter
                .emit_txn_for(job_duration, fullnode_txn_job)
                .await
                .map_err(|e| anyhow::format_err!("Failed to upgrade full node images: {}", e))?;
        }

        Ok(())
    }

    fn deadline(&self) -> Duration {
        Duration::from_secs(16 * 60)
    }
}

impl fmt::Display for CompatibilityTest {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "Compatibility test, phased upgrade to {} in batches of 1, {}, {}",
            self.updated_image_tag,
            self.first_batch.len(),
            self.second_batch.len()
        )
    }
}
<commit_msg>[cluster-test] Update compatibility tests for full nodes.<commit_after>// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

#![forbid(unsafe_code)]

use crate::{
    cluster::Cluster,
    experiments::{Context, Experiment, ExperimentParam},
    instance,
    instance::Instance,
    tx_emitter::EmitJobRequest,
};
use async_trait::async_trait;
use diem_logger::prelude::*;
use futures::future::try_join_all;
use std::{
    collections::HashSet,
    env, fmt,
    iter::{once, Iterator},
    time::{Duration, Instant},
};
use structopt::StructOpt;
use tokio::time;

/// Reboot `updated_instance` with newer image tag
pub async fn update_batch_instance(
    context: &mut Context<'_>,
    updated_instance: &[Instance],
    updated_lsr: &[Instance],
    updated_tag: String,
) -> anyhow::Result<()> {
    info!("Stop Existing instances.");
    let futures: Vec<_> = updated_instance.iter().map(Instance::stop).collect();
    try_join_all(futures).await?;

    if !updated_lsr.is_empty() {
        info!("Stop associated lsr instances.");
        let futures: Vec<_> = updated_lsr.iter().map(Instance::stop).collect();
        try_join_all(futures).await?;
        info!("Reinstantiate a set of new lsr.");
        let futures: Vec<_> = updated_lsr
            .iter()
            .map(|instance| {
                let mut newer_config = instance.instance_config().clone();
                newer_config.replace_tag(updated_tag.clone()).unwrap();
                context.cluster_swarm.spawn_new_instance(newer_config)
            })
            .collect();
        try_join_all(futures).await?;
        info!("Wait for the instance to sync up with peers");
        time::sleep(Duration::from_secs(20)).await;
    }

    info!("Reinstantiate a set of new nodes.");
    let futures: Vec<_> = updated_instance
        .iter()
        .map(|instance| {
            let mut newer_config = instance.instance_config().clone();
            newer_config.replace_tag(updated_tag.clone()).unwrap();
            context.cluster_swarm.spawn_new_instance(newer_config)
        })
        .collect();
    let instances = try_join_all(futures).await?;

    info!("Wait for the instances to recover.");
    let deadline = Instant::now() + Duration::from_secs(5 * 60);
    let futures: Vec<_> = instances
        .iter()
        .map(|instance| instance.wait_json_rpc(deadline))
        .collect();
    try_join_all(futures).await?;

    // Add a timeout to have wait for validators back to healthy mode.
    // TODO: Replace this with a blocking health check.
    info!("Wait for the instance to sync up with peers");
    time::sleep(Duration::from_secs(20)).await;
    Ok(())
}

pub fn get_corresponding_full_nodes<'a>(
    validators: impl Iterator<Item = &'a Instance>,
    cluster: &Cluster,
) -> Vec<Instance> {
    let validator_groups = validators
        .map(|instance| instance.validator_group())
        .collect::<Vec<_>>();
    cluster
        .fullnode_instances()
        .iter()
        .filter(|full_node| validator_groups.contains(&full_node.validator_group()))
        .cloned()
        .collect()
}

pub fn get_instance_list_str(batch: &[Instance]) -> String {
    let mut nodes_list = String::from("");
    for instance in batch.iter() {
        nodes_list.push_str(&instance.to_string());
        nodes_list.push_str(", ")
    }
    nodes_list
}

#[derive(StructOpt, Debug)]
pub struct CompatiblityTestParams {
    #[structopt(
        long,
        default_value = "15",
        help = "Number of nodes to update in the first batch"
    )]
    pub count: usize,
    #[structopt(long, help = "Image tag of newer validator software")]
    pub updated_image_tag: String,
}

pub struct CompatibilityTest {
    first_node: Instance,
    first_lsr: Vec<Instance>,
    first_full_node: Vec<Instance>,
    first_batch: Vec<Instance>,
    first_batch_lsr: Vec<Instance>,
    first_full_nodes_batch: Vec<Instance>,
    second_batch: Vec<Instance>,
    second_batch_lsr: Vec<Instance>,
    second_full_nodes_batch: Vec<Instance>,
    updated_image_tag: String,
}

impl ExperimentParam for CompatiblityTestParams {
    type E = CompatibilityTest;
    fn build(self, cluster: &Cluster) -> Self::E {
        if self.count > cluster.validator_instances().len() || self.count == 0 {
            panic!(
                "Can not reboot {} validators in cluster with {} instances",
                self.count,
                cluster.validator_instances().len()
            );
        }
        let (first_batch, second_batch) = cluster.split_n_validators_random(self.count);
        let mut first_batch = first_batch.into_validator_instances();
        let second_batch = second_batch.into_validator_instances();
        let first_node = first_batch
            .pop()
            .expect("Requires at least one validator in the first batch");
        let first_full_node = get_corresponding_full_nodes([first_node.clone()].iter(), cluster);

        let mut first_lsr = vec![];
        let mut first_batch_lsr = vec![];
        let mut second_batch_lsr = vec![];
        if !cluster.lsr_instances().is_empty() {
            first_batch_lsr = cluster.lsr_instances_for_validators(&first_batch);
            second_batch_lsr = cluster.lsr_instances_for_validators(&second_batch);
            first_lsr = cluster.lsr_instances_for_validators(&[first_node.clone()]);
        }

        let first_full_nodes_batch = get_corresponding_full_nodes(first_batch.iter(), cluster);
        let second_full_nodes_batch = get_corresponding_full_nodes(second_batch.iter(), cluster);

        Self::E {
            first_node,
            first_lsr,
            first_full_node,
            first_batch,
            first_batch_lsr,
            first_full_nodes_batch,
            second_batch,
            second_batch_lsr,
            second_full_nodes_batch,
            updated_image_tag: self.updated_image_tag,
        }
    }
}

#[async_trait]
impl Experiment for CompatibilityTest {
    fn affected_validators(&self) -> HashSet<String> {
        instance::instancelist_to_set(&self.first_batch)
            .union(&instance::instancelist_to_set(&self.second_batch))
            .cloned()
            .chain(once(self.first_node.peer_name().clone()))
            .collect()
    }

    async fn run(&mut self, context: &mut Context<'_>) -> anyhow::Result<()> {
        let job_duration = Duration::from_secs(3);
        context.report.report_text(format!(
            "Compatibility test results for {} ==> {} (PR)",
            context.current_tag, self.updated_image_tag
        ));

        // Generate some traffic
        let msg = format!(
            "1. All instances running {}, generating some traffic on network",
            context.current_tag
        );
        info!("{}", msg);
        context.report.report_text(msg);
        let all_full_nodes_request = EmitJobRequest::for_instances(
            context.cluster.validator_instances().to_vec(),
            context.global_emit_job_request,
            0,
            0,
        );
        context
            .tx_emitter
            .emit_txn_for(job_duration, all_full_nodes_request)
            .await
            .map_err(|e| anyhow::format_err!("Failed to generate traffic: {}", e))?;

        let msg = format!(
            "2. First full node {} ==> {}, to validate new full node to old validator node traffic",
            context.current_tag, self.updated_image_tag
        );
        info!("{}", msg);
        context.report.report_text(msg);
        info!(
            "Upgrade Full Node: {}",
            get_instance_list_str(&self.first_full_node)
        );
        update_batch_instance(
            context,
            &self.first_full_node,
            &[],
            self.updated_image_tag.clone(),
        )
        .await?;

        // Full node running at n+1, validator running n
        context
            .tx_emitter
            .emit_txn_for(
                job_duration,
                EmitJobRequest::for_instances(
                    self.first_full_node.clone(),
                    context.global_emit_job_request,
                    0,
                    0,
                ),
            )
            .await
            .map_err(|e| anyhow::format_err!("Storage backwards compat broken: {}", e))?;

        let msg = format!(
            "3. First Validator node {} ==> {}, to validate storage compatibility",
            context.current_tag, self.updated_image_tag
        );
        info!("{}", msg);
        context.report.report_text(msg);
        info!("Upgrading validator: {}", self.first_node);
        let first_node = vec![self.first_node.clone()];
        update_batch_instance(
            context,
            &first_node,
            &self.first_lsr,
            self.updated_image_tag.clone(),
        )
        .await?;
        context
            .tx_emitter
            .emit_txn_for(
                job_duration,
                EmitJobRequest::for_instances(
                    self.first_full_node.clone(),
                    context.global_emit_job_request,
                    0,
                    0,
                ),
            )
            .await
            .map_err(|e| anyhow::format_err!("Storage backwards compat broken: {}", e))?;

        let msg = format!(
            "4. First batch validators ({}) {} ==> {}, to test consensus and traffic between old full nodes and new validator node",
            self.first_batch.len(),
            context.current_tag,
            self.updated_image_tag
        );
        info!("{}", msg);
        info!(
            "Upgrading validators: {}",
            get_instance_list_str(&self.first_batch)
        );
        context.report.report_text(msg);
        update_batch_instance(
            context,
            &self.first_batch,
            &self.first_batch_lsr,
            self.updated_image_tag.clone(),
        )
        .await?;

        // Full node running at n, validator running n+1
        context
            .tx_emitter
            .emit_txn_for(
                job_duration,
                EmitJobRequest::for_instances(
                    self.first_full_nodes_batch.clone(),
                    context.global_emit_job_request,
                    0,
                    0,
                ),
            )
            .await
            .map_err(|e| anyhow::format_err!("Consensus backwards compat broken: {}", e))?;

        let msg = format!(
            "5. First batch full nodes ({}) {} ==> {}",
            self.first_full_nodes_batch.len(),
            context.current_tag,
            self.updated_image_tag
        );
        info!("{}", msg);
        context.report.report_text(msg);
        info!(
            "Upgrading full nodes: {}",
            get_instance_list_str(&self.first_full_nodes_batch)
        );
        update_batch_instance(
            context,
            &self.first_full_nodes_batch,
            &[],
            self.updated_image_tag.clone(),
        )
        .await?;

        let msg = format!(
            "6. Second batch validators ({}) {} ==> {}, to upgrade rest of the validators",
            self.second_batch.len(),
            context.current_tag,
            self.updated_image_tag
        );
        info!("{}", msg);
        context.report.report_text(msg);
        info!(
            "Upgrading validators: {}",
            get_instance_list_str(&self.second_batch)
        );
        update_batch_instance(
            context,
            &self.second_batch,
            &self.second_batch_lsr,
            self.updated_image_tag.clone(),
        )
        .await?;
        context
            .tx_emitter
            .emit_txn_for(
                job_duration,
                EmitJobRequest::for_instances(
                    self.second_batch.clone(),
                    context.global_emit_job_request,
                    0,
                    0,
                ),
            )
            .await
            .map_err(|e| {
                anyhow::format_err!("Failed to upgrade rest of validator images: {}", e)
            })?;

        let msg = format!(
            "7. Second batch of full nodes ({}) {} ==> {}, to finish the network upgrade",
            self.second_full_nodes_batch.len(),
            context.current_tag,
            self.updated_image_tag
        );
        info!("{}", msg);
        info!(
            "Upgrading full nodes: {}",
            get_instance_list_str(&self.second_full_nodes_batch)
        );
        context.report.report_text(msg);
        update_batch_instance(
            context,
            &self.second_full_nodes_batch,
            &[],
            self.updated_image_tag.clone(),
        )
        .await?;
        context
            .tx_emitter
            .emit_txn_for(
                job_duration,
                EmitJobRequest::for_instances(
                    self.second_full_nodes_batch.clone(),
                    context.global_emit_job_request,
                    0,
                    0,
                ),
            )
            .await
            .map_err(|e| anyhow::format_err!("Failed to upgrade full node images: {}", e))?;

        Ok(())
    }

    fn deadline(&self) -> Duration {
        Duration::from_secs(16 * 60)
    }
}

impl fmt::Display for CompatibilityTest {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "Compatibility test, phased upgrade to {} in batches of 1, {}, {}",
            self.updated_image_tag,
            self.first_batch.len(),
            self.second_batch.len()
        )
    }
}
<|endoftext|>


<commit_before>use crate::{Assist, AssistId, AssistCtx};

use hir::Resolver;
use hir::db::HirDatabase;
use ra_syntax::{SmolStr, SyntaxKind, TextRange, TextUnit, TreeArc};
use ra_syntax::ast::{self, AstNode, FnDef, ImplItem, ImplItemKind, NameOwner};
use ra_db::FilePosition;
use ra_fmt::{leading_indent, reindent};

use itertools::Itertools;

/// Given an `ast::ImplBlock`, resolves the target trait (the one being
/// implemented) to a `ast::TraitDef`.
fn resolve_target_trait_def(
    db: &impl HirDatabase,
    resolver: &Resolver,
    impl_block: &ast::ImplBlock,
) -> Option<TreeArc<ast::TraitDef>> {
    let ast_path = impl_block.target_trait().map(AstNode::syntax).and_then(ast::PathType::cast)?;
    let hir_path = ast_path.path().and_then(hir::Path::from_ast)?;

    match resolver.resolve_path(db, &hir_path).take_types() {
        Some(hir::Resolution::Def(hir::ModuleDef::Trait(def))) => Some(def.source(db).1),
        _ => None,
    }
}

fn build_func_body(def: &ast::FnDef) -> String {
    let mut buf = String::new();

    for child in def.syntax().children() {
        if child.kind() == SyntaxKind::SEMI {
            buf.push_str(" { unimplemented!() }")
        } else {
            child.text().push_to(&mut buf);
        }
    }

    buf.trim_end().to_string()
}

pub(crate) fn add_missing_impl_members(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let impl_node = ctx.node_at_offset::<ast::ImplBlock>()?;
    let impl_item_list = impl_node.item_list()?;

    let trait_def = {
        let file_id = ctx.frange.file_id;
        let position = FilePosition { file_id, offset: impl_node.syntax().range().start() };
        let resolver = hir::source_binder::resolver_for_position(ctx.db, position);

        resolve_target_trait_def(ctx.db, &resolver, impl_node)?
    };

    let missing_fns: Vec<_> = {
        let fn_def_opt = |kind| if let ImplItemKind::FnDef(def) = kind { Some(def) } else { None };
        let def_name = |def| -> Option<&SmolStr> { FnDef::name(def).map(ast::Name::text) };

        let trait_items =
            trait_def.syntax().descendants().find_map(ast::ItemList::cast)?.impl_items();
        let impl_items = impl_item_list.impl_items();

        let trait_fns = trait_items.map(ImplItem::kind).filter_map(fn_def_opt).collect::<Vec<_>>();
        let impl_fns = impl_items.map(ImplItem::kind).filter_map(fn_def_opt).collect::<Vec<_>>();

        trait_fns
            .into_iter()
            .filter(|t| def_name(t).is_some())
            .filter(|t| impl_fns.iter().all(|i| def_name(i) != def_name(t)))
            .collect()
    };
    if missing_fns.is_empty() {
        return None;
    }

    ctx.add_action(AssistId("add_impl_missing_members"), "add missing impl members", |edit| {
        let (parent_indent, indent) = {
            // FIXME: Find a way to get the indent already used in the file.
            // Now, we copy the indent of first item or indent with 4 spaces relative to impl block
            const DEFAULT_INDENT: &str = "    ";
            let first_item = impl_item_list.impl_items().next();
            let first_item_indent =
                first_item.and_then(|i| leading_indent(i.syntax())).map(ToOwned::to_owned);
            let impl_block_indent = leading_indent(impl_node.syntax()).unwrap_or_default();

            (
                impl_block_indent.to_owned(),
                first_item_indent.unwrap_or_else(|| impl_block_indent.to_owned() + DEFAULT_INDENT),
            )
        };

        let changed_range = {
            let children = impl_item_list.syntax().children();
            let last_whitespace = children.filter_map(ast::Whitespace::cast).last();

            last_whitespace.map(|w| w.syntax().range()).unwrap_or_else(|| {
                let in_brackets = impl_item_list.syntax().range().end() - TextUnit::of_str("}");
                TextRange::from_to(in_brackets, in_brackets)
            })
        };

        let func_bodies = format!("\n{}", missing_fns.into_iter().map(build_func_body).join("\n"));
        let trailing_whitespace = format!("\n{}", parent_indent);
        let func_bodies = reindent(&func_bodies, &indent) + &trailing_whitespace;

        let replaced_text_range = TextUnit::of_str(&func_bodies);

        edit.replace(changed_range, func_bodies);
        edit.set_cursor(
            changed_range.start() + replaced_text_range - TextUnit::of_str(&trailing_whitespace),
        );
    });

    ctx.build()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::helpers::{check_assist, check_assist_not_applicable};

    #[test]
    fn test_add_missing_impl_members() {
        check_assist(
            add_missing_impl_members,
            "
trait Foo {
    fn foo(&self);
    fn bar(&self);
    fn baz(&self);
}

struct S;

impl Foo for S {
    fn bar(&self) {}
    <|>
}",
            "
trait Foo {
    fn foo(&self);
    fn bar(&self);
    fn baz(&self);
}

struct S;

impl Foo for S {
    fn bar(&self) {}
    fn foo(&self) { unimplemented!() }
    fn baz(&self) { unimplemented!() }<|>
}",
        );
    }

    #[test]
    fn test_copied_overriden_members() {
        check_assist(
            add_missing_impl_members,
            "
trait Foo {
    fn foo(&self);
    fn bar(&self) -> bool { true }
    fn baz(&self) -> u32 { 42 }
}

struct S;

impl Foo for S {
    fn bar(&self) {}
    <|>
}",
            "
trait Foo {
    fn foo(&self);
    fn bar(&self) -> bool { true }
    fn baz(&self) -> u32 { 42 }
}

struct S;

impl Foo for S {
    fn bar(&self) {}
    fn foo(&self) { unimplemented!() }
    fn baz(&self) -> u32 { 42 }<|>
}",
        );
    }

    #[test]
    fn test_empty_impl_block() {
        check_assist(
            add_missing_impl_members,
            "
trait Foo { fn foo(&self); }
struct S;
impl Foo for S {<|>}",
            "
trait Foo { fn foo(&self); }
struct S;
impl Foo for S {
    fn foo(&self) { unimplemented!() }<|>
}",
        );
    }

    #[test]
    fn test_cursor_after_empty_impl_block() {
        check_assist(
            add_missing_impl_members,
            "
trait Foo { fn foo(&self); }
struct S;
impl Foo for S {}<|>",
            "
trait Foo { fn foo(&self); }
struct S;
impl Foo for S {
    fn foo(&self) { unimplemented!() }<|>
}",
        )
    }

    #[test]
    fn test_empty_trait() {
        check_assist_not_applicable(
            add_missing_impl_members,
            "
trait Foo;
struct S;
impl Foo for S { <|> }",
        )
    }

    #[test]
    fn test_ignore_unnamed_trait_members() {
        check_assist(
            add_missing_impl_members,
            "
trait Foo {
    fn (arg: u32);
    fn valid(some: u32) -> bool { false }
}
struct S;
impl Foo for S { <|> }",
            "
trait Foo {
    fn (arg: u32);
    fn valid(some: u32) -> bool { false }
}
struct S;
impl Foo for S {
    fn valid(some: u32) -> bool { false }<|>
}",
        )
    }

    #[test]
    fn test_indented_impl_block() {
        check_assist(
            add_missing_impl_members,
            "
trait Foo {
    fn valid(some: u32) -> bool { false }
}
struct S;

mod my_mod {
    impl crate::Foo for S { <|> }
}",
            "
trait Foo {
    fn valid(some: u32) -> bool { false }
}
struct S;

mod my_mod {
    impl crate::Foo for S {
        fn valid(some: u32) -> bool { false }<|>
    }
}",
        )
    }
}
<commit_msg>Move the primary assist fn to the top of the file<commit_after>use crate::{Assist, AssistId, AssistCtx};

use hir::Resolver;
use hir::db::HirDatabase;
use ra_syntax::{SmolStr, SyntaxKind, TextRange, TextUnit, TreeArc};
use ra_syntax::ast::{self, AstNode, FnDef, ImplItem, ImplItemKind, NameOwner};
use ra_db::FilePosition;
use ra_fmt::{leading_indent, reindent};

use itertools::Itertools;

pub(crate) fn add_missing_impl_members(mut ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    let impl_node = ctx.node_at_offset::<ast::ImplBlock>()?;
    let impl_item_list = impl_node.item_list()?;

    let trait_def = {
        let file_id = ctx.frange.file_id;
        let position = FilePosition { file_id, offset: impl_node.syntax().range().start() };
        let resolver = hir::source_binder::resolver_for_position(ctx.db, position);

        resolve_target_trait_def(ctx.db, &resolver, impl_node)?
    };

    let missing_fns: Vec<_> = {
        let fn_def_opt = |kind| if let ImplItemKind::FnDef(def) = kind { Some(def) } else { None };
        let def_name = |def| -> Option<&SmolStr> { FnDef::name(def).map(ast::Name::text) };

        let trait_items =
            trait_def.syntax().descendants().find_map(ast::ItemList::cast)?.impl_items();
        let impl_items = impl_item_list.impl_items();

        let trait_fns = trait_items.map(ImplItem::kind).filter_map(fn_def_opt).collect::<Vec<_>>();
        let impl_fns = impl_items.map(ImplItem::kind).filter_map(fn_def_opt).collect::<Vec<_>>();

        trait_fns
            .into_iter()
            .filter(|t| def_name(t).is_some())
            .filter(|t| impl_fns.iter().all(|i| def_name(i) != def_name(t)))
            .collect()
    };
    if missing_fns.is_empty() {
        return None;
    }

    ctx.add_action(AssistId("add_impl_missing_members"), "add missing impl members", |edit| {
        let (parent_indent, indent) = {
            // FIXME: Find a way to get the indent already used in the file.
            // Now, we copy the indent of first item or indent with 4 spaces relative to impl block
            const DEFAULT_INDENT: &str = "    ";
            let first_item = impl_item_list.impl_items().next();
            let first_item_indent =
                first_item.and_then(|i| leading_indent(i.syntax())).map(ToOwned::to_owned);
            let impl_block_indent = leading_indent(impl_node.syntax()).unwrap_or_default();

            (
                impl_block_indent.to_owned(),
                first_item_indent.unwrap_or_else(|| impl_block_indent.to_owned() + DEFAULT_INDENT),
            )
        };

        let changed_range = {
            let children = impl_item_list.syntax().children();
            let last_whitespace = children.filter_map(ast::Whitespace::cast).last();

            last_whitespace.map(|w| w.syntax().range()).unwrap_or_else(|| {
                let in_brackets = impl_item_list.syntax().range().end() - TextUnit::of_str("}");
                TextRange::from_to(in_brackets, in_brackets)
            })
        };

        let func_bodies = format!("\n{}", missing_fns.into_iter().map(build_func_body).join("\n"));
        let trailing_whitespace = format!("\n{}", parent_indent);
        let func_bodies = reindent(&func_bodies, &indent) + &trailing_whitespace;

        let replaced_text_range = TextUnit::of_str(&func_bodies);

        edit.replace(changed_range, func_bodies);
        edit.set_cursor(
            changed_range.start() + replaced_text_range - TextUnit::of_str(&trailing_whitespace),
        );
    });

    ctx.build()
}

/// Given an `ast::ImplBlock`, resolves the target trait (the one being
/// implemented) to a `ast::TraitDef`.
fn resolve_target_trait_def(
    db: &impl HirDatabase,
    resolver: &Resolver,
    impl_block: &ast::ImplBlock,
) -> Option<TreeArc<ast::TraitDef>> {
    let ast_path = impl_block.target_trait().map(AstNode::syntax).and_then(ast::PathType::cast)?;
    let hir_path = ast_path.path().and_then(hir::Path::from_ast)?;

    match resolver.resolve_path(db, &hir_path).take_types() {
        Some(hir::Resolution::Def(hir::ModuleDef::Trait(def))) => Some(def.source(db).1),
        _ => None,
    }
}

fn build_func_body(def: &ast::FnDef) -> String {
    let mut buf = String::new();

    for child in def.syntax().children() {
        if child.kind() == SyntaxKind::SEMI {
            buf.push_str(" { unimplemented!() }")
        } else {
            child.text().push_to(&mut buf);
        }
    }

    buf.trim_end().to_string()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::helpers::{check_assist, check_assist_not_applicable};

    #[test]
    fn test_add_missing_impl_members() {
        check_assist(
            add_missing_impl_members,
            "
trait Foo {
    fn foo(&self);
    fn bar(&self);
    fn baz(&self);
}

struct S;

impl Foo for S {
    fn bar(&self) {}
    <|>
}",
            "
trait Foo {
    fn foo(&self);
    fn bar(&self);
    fn baz(&self);
}

struct S;

impl Foo for S {
    fn bar(&self) {}
    fn foo(&self) { unimplemented!() }
    fn baz(&self) { unimplemented!() }<|>
}",
        );
    }

    #[test]
    fn test_copied_overriden_members() {
        check_assist(
            add_missing_impl_members,
            "
trait Foo {
    fn foo(&self);
    fn bar(&self) -> bool { true }
    fn baz(&self) -> u32 { 42 }
}

struct S;

impl Foo for S {
    fn bar(&self) {}
    <|>
}",
            "
trait Foo {
    fn foo(&self);
    fn bar(&self) -> bool { true }
    fn baz(&self) -> u32 { 42 }
}

struct S;

impl Foo for S {
    fn bar(&self) {}
    fn foo(&self) { unimplemented!() }
    fn baz(&self) -> u32 { 42 }<|>
}",
        );
    }

    #[test]
    fn test_empty_impl_block() {
        check_assist(
            add_missing_impl_members,
            "
trait Foo { fn foo(&self); }
struct S;
impl Foo for S {<|>}",
            "
trait Foo { fn foo(&self); }
struct S;
impl Foo for S {
    fn foo(&self) { unimplemented!() }<|>
}",
        );
    }

    #[test]
    fn test_cursor_after_empty_impl_block() {
        check_assist(
            add_missing_impl_members,
            "
trait Foo { fn foo(&self); }
struct S;
impl Foo for S {}<|>",
            "
trait Foo { fn foo(&self); }
struct S;
impl Foo for S {
    fn foo(&self) { unimplemented!() }<|>
}",
        )
    }

    #[test]
    fn test_empty_trait() {
        check_assist_not_applicable(
            add_missing_impl_members,
            "
trait Foo;
struct S;
impl Foo for S { <|> }",
        )
    }

    #[test]
    fn test_ignore_unnamed_trait_members() {
        check_assist(
            add_missing_impl_members,
            "
trait Foo {
    fn (arg: u32);
    fn valid(some: u32) -> bool { false }
}
struct S;
impl Foo for S { <|> }",
            "
trait Foo {
    fn (arg: u32);
    fn valid(some: u32) -> bool { false }
}
struct S;
impl Foo for S {
    fn valid(some: u32) -> bool { false }<|>
}",
        )
    }

    #[test]
    fn test_indented_impl_block() {
        check_assist(
            add_missing_impl_members,
            "
trait Foo {
    fn valid(some: u32) -> bool { false }
}
struct S;

mod my_mod {
    impl crate::Foo for S { <|> }
}",
            "
trait Foo {
    fn valid(some: u32) -> bool { false }
}
struct S;

mod my_mod {
    impl crate::Foo for S {
        fn valid(some: u32) -> bool { false }<|>
    }
}",
        )
    }
}
<|endoftext|>


<commit_before>//! FIXME: write short doc here

use ra_syntax::{
    ast::{self, ModuleItemOwner},
    match_ast, AstNode,
};
use rustc_hash::FxHashMap;

use crate::completion::{CompletionContext, CompletionItem, CompletionKind, Completions};

/// Complete repeated parameters, both name and type. For example, if all
/// functions in a file have a `spam: &mut Spam` parameter, a completion with
/// `spam: &mut Spam` insert text/label and `spam` lookup string will be
/// suggested.
pub(super) fn complete_fn_param(acc: &mut Completions, ctx: &CompletionContext) {
    if !ctx.is_param {
        return;
    }

    let mut params = FxHashMap::default();
    for node in ctx.token.parent().ancestors() {
        let items = match_ast! {
            match node {
                ast::SourceFile(it) => it.items(),
                ast::ItemList(it) => it.items(),
                _ => continue,
            }
        };
        for item in items {
            if let ast::ModuleItem::FnDef(func) = item {
                func.param_list().into_iter().flat_map(|it| it.params()).for_each(|param| {
                    let text = param.syntax().text().to_string();
                    params.entry(text).or_insert((0, param)).0 += 1;
                })
            }
        }
    }
    params
        .into_iter()
        .filter_map(|(label, (count, param))| {
            let lookup = param.pat()?.syntax().text().to_string();
            if count < 2 {
                None
            } else {
                Some((label, lookup))
            }
        })
        .for_each(|(label, lookup)| {
            CompletionItem::new(CompletionKind::Magic, ctx.source_range(), label)
                .lookup_by(lookup)
                .add_to(acc)
        });
}

#[cfg(test)]
mod tests {
    use expect::{expect, Expect};

    use crate::completion::{test_utils::do_completion, CompletionKind};

    fn check(ra_fixture: &str, expect: Expect) {
        let actual = do_completion(ra_fixture, CompletionKind::Magic);
        expect.assert_debug_eq(&actual);
    }

    #[test]
    fn test_param_completion_last_param() {
        check(
            r#"
fn foo(file_id: FileId) {}
fn bar(file_id: FileId) {}
fn baz(file<|>) {}
"#,
            expect![[r#"
                [
                    CompletionItem {
                        label: "file_id: FileId",
                        source_range: 61..65,
                        delete: 61..65,
                        insert: "file_id: FileId",
                        lookup: "file_id",
                    },
                ]
            "#]],
        );
    }

    #[test]
    fn test_param_completion_nth_param() {
        check(
            r#"
fn foo(file_id: FileId) {}
fn bar(file_id: FileId) {}
fn baz(file<|>, x: i32) {}
"#,
            expect![[r#"
                [
                    CompletionItem {
                        label: "file_id: FileId",
                        source_range: 61..65,
                        delete: 61..65,
                        insert: "file_id: FileId",
                        lookup: "file_id",
                    },
                ]
            "#]],
        );
    }

    #[test]
    fn test_param_completion_trait_param() {
        check(
            r#"
pub(crate) trait SourceRoot {
    pub fn contains(&self, file_id: FileId) -> bool;
    pub fn module_map(&self) -> &ModuleMap;
    pub fn lines(&self, file_id: FileId) -> &LineIndex;
    pub fn syntax(&self, file<|>)
}
"#,
            expect![[r#"
                [
                    CompletionItem {
                        label: "file_id: FileId",
                        source_range: 208..212,
                        delete: 208..212,
                        insert: "file_id: FileId",
                        lookup: "file_id",
                    },
                ]
            "#]],
        );
    }
}
<commit_msg>Refresh tests<commit_after>//! FIXME: write short doc here

use ra_syntax::{
    ast::{self, ModuleItemOwner},
    match_ast, AstNode,
};
use rustc_hash::FxHashMap;

use crate::completion::{CompletionContext, CompletionItem, CompletionKind, Completions};

/// Complete repeated parameters, both name and type. For example, if all
/// functions in a file have a `spam: &mut Spam` parameter, a completion with
/// `spam: &mut Spam` insert text/label and `spam` lookup string will be
/// suggested.
pub(super) fn complete_fn_param(acc: &mut Completions, ctx: &CompletionContext) {
    if !ctx.is_param {
        return;
    }

    let mut params = FxHashMap::default();
    for node in ctx.token.parent().ancestors() {
        let items = match_ast! {
            match node {
                ast::SourceFile(it) => it.items(),
                ast::ItemList(it) => it.items(),
                _ => continue,
            }
        };
        for item in items {
            if let ast::ModuleItem::FnDef(func) = item {
                func.param_list().into_iter().flat_map(|it| it.params()).for_each(|param| {
                    let text = param.syntax().text().to_string();
                    params.entry(text).or_insert((0, param)).0 += 1;
                })
            }
        }
    }
    params
        .into_iter()
        .filter_map(|(label, (count, param))| {
            let lookup = param.pat()?.syntax().text().to_string();
            if count < 2 {
                None
            } else {
                Some((label, lookup))
            }
        })
        .for_each(|(label, lookup)| {
            CompletionItem::new(CompletionKind::Magic, ctx.source_range(), label)
                .kind(crate::CompletionItemKind::Binding)
                .lookup_by(lookup)
                .add_to(acc)
        });
}

#[cfg(test)]
mod tests {
    use expect::{expect, Expect};

    use crate::completion::{test_utils::completion_list, CompletionKind};

    fn check(ra_fixture: &str, expect: Expect) {
        let actual = completion_list(ra_fixture, CompletionKind::Magic);
        expect.assert_eq(&actual);
    }

    #[test]
    fn test_param_completion_last_param() {
        check(
            r#"
fn foo(file_id: FileId) {}
fn bar(file_id: FileId) {}
fn baz(file<|>) {}
"#,
            expect![[r#"
                bn file_id: FileId
            "#]],
        );
    }

    #[test]
    fn test_param_completion_nth_param() {
        check(
            r#"
fn foo(file_id: FileId) {}
fn bar(file_id: FileId) {}
fn baz(file<|>, x: i32) {}
"#,
            expect![[r#"
                bn file_id: FileId
            "#]],
        );
    }

    #[test]
    fn test_param_completion_trait_param() {
        check(
            r#"
pub(crate) trait SourceRoot {
    pub fn contains(&self, file_id: FileId) -> bool;
    pub fn module_map(&self) -> &ModuleMap;
    pub fn lines(&self, file_id: FileId) -> &LineIndex;
    pub fn syntax(&self, file<|>)
}
"#,
            expect![[r#"
                bn file_id: FileId
            "#]],
        );
    }
}
<|endoftext|>


<commit_before>use super::*;
use int_to_bytes::{int_to_bytes32, int_to_bytes8};

#[derive(Clone)]
pub struct Inner {
    pub a: u64,
    pub b: u64,
    pub c: u64,
    pub d: u64,
}

impl CachedTreeHash for Inner {
    type Item = Self;

    fn build_cache_bytes(&self) -> Vec<u8> {
        let mut leaves = vec![];

        leaves.append(&mut self.a.build_cache_bytes());
        leaves.append(&mut self.b.build_cache_bytes());
        leaves.append(&mut self.c.build_cache_bytes());
        leaves.append(&mut self.d.build_cache_bytes());

        merkleize(leaves)
    }

    fn num_bytes(&self) -> usize {
        let mut bytes = 0;
        bytes += self.a.num_bytes();
        bytes += self.b.num_bytes();
        bytes += self.c.num_bytes();
        bytes += self.d.num_bytes();
        bytes
    }

    fn cached_hash_tree_root(
        &self,
        other: &Self,
        cache: &mut TreeHashCache,
        chunk: usize,
    ) -> Option<usize> {
        let mut num_leaves: usize = 0;
        num_leaves += num_unsanitized_leaves(self.a.num_bytes());
        num_leaves += num_unsanitized_leaves(self.b.num_bytes());
        num_leaves += num_unsanitized_leaves(self.c.num_bytes());
        num_leaves += num_unsanitized_leaves(self.d.num_bytes());

        let num_nodes = num_nodes(num_leaves);
        let num_internal_nodes = num_nodes - num_leaves;

        // Skip past the internal nodes and update any changed leaf nodes.
        {
            let chunk = chunk + num_internal_nodes;
            let chunk = self.a.cached_hash_tree_root(&other.a, cache, chunk)?;
            let chunk = self.b.cached_hash_tree_root(&other.b, cache, chunk)?;
            let chunk = self.c.cached_hash_tree_root(&other.c, cache, chunk)?;
            let _chunk = self.d.cached_hash_tree_root(&other.d, cache, chunk)?;
        }

        // Iterate backwards through the internal nodes, rehashing any node where it's children
        // have changed.
        for chunk in (chunk..chunk + num_internal_nodes).into_iter().rev() {
            if cache.children_modified(chunk)? {
                cache.modify_chunk(chunk, &cache.hash_children(chunk)?)?;
            }
        }

        Some(chunk + num_nodes)
    }
}

fn join(many: Vec<Vec<u8>>) -> Vec<u8> {
    let mut all = vec![];
    for one in many {
        all.extend_from_slice(&mut one.clone())
    }
    all
}

#[test]
fn partial_modification_u64_vec() {
    let n: u64 = 50;

    let original_vec: Vec<u64> = (0..n).collect();

    // Generate initial cache.
    let original_cache = original_vec.build_cache_bytes();

    // Modify the vec
    let mut modified_vec = original_vec.clone();
    modified_vec[n as usize - 1] = 42;

    // Perform a differential hash
    let mut cache_struct = TreeHashCache::from_bytes(original_cache.clone()).unwrap();
    modified_vec.cached_hash_tree_root(&original_vec, &mut cache_struct, 0);
    let modified_cache: Vec<u8> = cache_struct.into();

    // Generate reference data.
    let mut data = vec![];
    for i in &modified_vec {
        data.append(&mut int_to_bytes8(*i));
    }
    let data = sanitise_bytes(data);
    let expected = merkleize(data);

    assert_eq!(expected, modified_cache);
}

#[test]
fn large_vec_of_u64_builds() {
    let n: u64 = 50;

    let my_vec: Vec<u64> = (0..n).collect();

    // Generate function output.
    let cache = my_vec.build_cache_bytes();

    // Generate reference data.
    let mut data = vec![];
    for i in &my_vec {
        data.append(&mut int_to_bytes8(*i));
    }
    let data = sanitise_bytes(data);
    let expected = merkleize(data);

    assert_eq!(expected, cache);
}

#[test]
fn vec_of_u64_builds() {
    let data = join(vec![
        int_to_bytes8(1),
        int_to_bytes8(2),
        int_to_bytes8(3),
        int_to_bytes8(4),
        int_to_bytes8(5),
        vec![0; 32 - 8], // padding
    ]);

    let expected = merkleize(data);

    let my_vec = vec![1, 2, 3, 4, 5];

    let cache = my_vec.build_cache_bytes();

    assert_eq!(expected, cache);
}

#[test]
fn merkleize_odd() {
    let data = join(vec![
        int_to_bytes32(1),
        int_to_bytes32(2),
        int_to_bytes32(3),
        int_to_bytes32(4),
        int_to_bytes32(5),
    ]);

    merkleize(sanitise_bytes(data));
}

fn generic_test(index: usize) {
    let inner = Inner {
        a: 1,
        b: 2,
        c: 3,
        d: 4,
    };

    let cache = inner.build_cache_bytes();

    let changed_inner = match index {
        0 => Inner {
            a: 42,
            ..inner.clone()
        },
        1 => Inner {
            b: 42,
            ..inner.clone()
        },
        2 => Inner {
            c: 42,
            ..inner.clone()
        },
        3 => Inner {
            d: 42,
            ..inner.clone()
        },
        _ => panic!("bad index"),
    };

    let mut cache_struct = TreeHashCache::from_bytes(cache.clone()).unwrap();

    changed_inner
        .cached_hash_tree_root(&inner, &mut cache_struct, 0)
        .unwrap();

    // assert_eq!(*cache_struct.hash_count, 3);

    let new_cache: Vec<u8> = cache_struct.into();

    let data1 = int_to_bytes32(1);
    let data2 = int_to_bytes32(2);
    let data3 = int_to_bytes32(3);
    let data4 = int_to_bytes32(4);

    let mut data = vec![data1, data2, data3, data4];

    data[index] = int_to_bytes32(42);

    let expected = merkleize(join(data));

    assert_eq!(expected, new_cache);
}

#[test]
fn cached_hash_on_inner() {
    generic_test(0);
    generic_test(1);
    generic_test(2);
    generic_test(3);
}

#[test]
fn build_cache_matches_merkelize() {
    let data1 = int_to_bytes32(1);
    let data2 = int_to_bytes32(2);
    let data3 = int_to_bytes32(3);
    let data4 = int_to_bytes32(4);

    let data = join(vec![data1, data2, data3, data4]);
    let expected = merkleize(data);

    let inner = Inner {
        a: 1,
        b: 2,
        c: 3,
        d: 4,
    };

    let cache = inner.build_cache_bytes();

    assert_eq!(expected, cache);
}

#[test]
fn merkleize_4_leaves() {
    let data1 = hash(&int_to_bytes32(1));
    let data2 = hash(&int_to_bytes32(2));
    let data3 = hash(&int_to_bytes32(3));
    let data4 = hash(&int_to_bytes32(4));

    let data = join(vec![
        data1.clone(),
        data2.clone(),
        data3.clone(),
        data4.clone(),
    ]);

    let cache = merkleize(data);

    let hash_12 = {
        let mut joined = vec![];
        joined.append(&mut data1.clone());
        joined.append(&mut data2.clone());
        hash(&joined)
    };
    let hash_34 = {
        let mut joined = vec![];
        joined.append(&mut data3.clone());
        joined.append(&mut data4.clone());
        hash(&joined)
    };
    let hash_hash12_hash_34 = {
        let mut joined = vec![];
        joined.append(&mut hash_12.clone());
        joined.append(&mut hash_34.clone());
        hash(&joined)
    };

    for (i, chunk) in cache.chunks(HASHSIZE).enumerate().rev() {
        let expected = match i {
            0 => hash_hash12_hash_34.clone(),
            1 => hash_12.clone(),
            2 => hash_34.clone(),
            3 => data1.clone(),
            4 => data2.clone(),
            5 => data3.clone(),
            6 => data4.clone(),
            _ => vec![],
        };

        assert_eq!(chunk, &expected[..], "failed at {}", i);
    }
}
<commit_msg>Build breaking recursion tests for cache hashing<commit_after>use super::*;
use int_to_bytes::{int_to_bytes32, int_to_bytes8};

#[derive(Clone)]
pub struct Inner {
    pub a: u64,
    pub b: u64,
    pub c: u64,
    pub d: u64,
}

impl CachedTreeHash for Inner {
    type Item = Self;

    fn build_cache_bytes(&self) -> Vec<u8> {
        let mut leaves = vec![];

        leaves.append(&mut self.a.build_cache_bytes());
        leaves.append(&mut self.b.build_cache_bytes());
        leaves.append(&mut self.c.build_cache_bytes());
        leaves.append(&mut self.d.build_cache_bytes());

        merkleize(leaves)
    }

    fn num_bytes(&self) -> usize {
        let mut bytes = 0;
        bytes += self.a.num_bytes();
        bytes += self.b.num_bytes();
        bytes += self.c.num_bytes();
        bytes += self.d.num_bytes();
        bytes
    }

    fn cached_hash_tree_root(
        &self,
        other: &Self,
        cache: &mut TreeHashCache,
        chunk: usize,
    ) -> Option<usize> {
        let mut num_leaves: usize = 0;
        num_leaves += num_unsanitized_leaves(self.a.num_bytes());
        num_leaves += num_unsanitized_leaves(self.b.num_bytes());
        num_leaves += num_unsanitized_leaves(self.c.num_bytes());
        num_leaves += num_unsanitized_leaves(self.d.num_bytes());

        let num_nodes = num_nodes(num_leaves);
        let num_internal_nodes = num_nodes - num_leaves;

        // Skip past the internal nodes and update any changed leaf nodes.
        {
            let chunk = chunk + num_internal_nodes;
            let chunk = self.a.cached_hash_tree_root(&other.a, cache, chunk)?;
            let chunk = self.b.cached_hash_tree_root(&other.b, cache, chunk)?;
            let chunk = self.c.cached_hash_tree_root(&other.c, cache, chunk)?;
            let _chunk = self.d.cached_hash_tree_root(&other.d, cache, chunk)?;
        }

        // Iterate backwards through the internal nodes, rehashing any node where it's children
        // have changed.
        for chunk in (chunk..chunk + num_internal_nodes).into_iter().rev() {
            if cache.children_modified(chunk)? {
                cache.modify_chunk(chunk, &cache.hash_children(chunk)?)?;
            }
        }

        Some(chunk + num_nodes)
    }
}

#[derive(Clone)]
pub struct Outer {
    pub a: u64,
    pub b: Inner,
    pub c: u64,
}

impl CachedTreeHash for Outer {
    type Item = Self;

    fn build_cache_bytes(&self) -> Vec<u8> {
        let mut leaves = vec![];

        leaves.append(&mut self.a.build_cache_bytes());
        leaves.append(&mut self.b.build_cache_bytes());
        leaves.append(&mut self.c.build_cache_bytes());

        merkleize(leaves)
    }

    fn num_bytes(&self) -> usize {
        let mut bytes = 0;
        bytes += self.a.num_bytes();
        bytes += self.b.num_bytes();
        bytes
    }

    fn cached_hash_tree_root(
        &self,
        other: &Self,
        cache: &mut TreeHashCache,
        chunk: usize,
    ) -> Option<usize> {
        let mut num_leaves: usize = 0;
        num_leaves += num_unsanitized_leaves(self.a.num_bytes());
        num_leaves += num_unsanitized_leaves(self.b.num_bytes());
        num_leaves += num_unsanitized_leaves(self.c.num_bytes());

        let num_nodes = num_nodes(num_leaves);
        let num_internal_nodes = num_nodes - num_leaves;

        // Skip past the internal nodes and update any changed leaf nodes.
        {
            let chunk = chunk + num_internal_nodes;
            let chunk = self.a.cached_hash_tree_root(&other.a, cache, chunk)?;
            let chunk = self.b.cached_hash_tree_root(&other.b, cache, chunk)?;
            let _chunk = self.c.cached_hash_tree_root(&other.c, cache, chunk)?;
        }

        // Iterate backwards through the internal nodes, rehashing any node where it's children
        // have changed.
        for chunk in (chunk..chunk + num_internal_nodes).into_iter().rev() {
            if cache.children_modified(chunk)? {
                cache.modify_chunk(chunk, &cache.hash_children(chunk)?)?;
            }
        }

        Some(chunk + num_nodes)
    }
}

fn join(many: Vec<Vec<u8>>) -> Vec<u8> {
    let mut all = vec![];
    for one in many {
        all.extend_from_slice(&mut one.clone())
    }
    all
}

#[test]
fn partial_modification_to_outer() {
    let inner = Inner {
        a: 1,
        b: 2,
        c: 3,
        d: 4,
    };

    let original_outer = Outer {
        a: 0,
        b: inner.clone(),
        c: 5,
    };

    // Build the initial cache.
    let original_cache = original_outer.build_cache_bytes();

    // Modify outer
    let modified_outer = Outer {
        c: 42,
        ..original_outer.clone()
    };

    // Perform a differential hash
    let mut cache_struct = TreeHashCache::from_bytes(original_cache.clone()).unwrap();
    modified_outer.cached_hash_tree_root(&original_outer, &mut cache_struct, 0);
    let modified_cache: Vec<u8> = cache_struct.into();

    // Generate reference data.
    let mut data = vec![];
    data.append(&mut int_to_bytes32(0));
    data.append(&mut inner.build_cache_bytes());
    data.append(&mut int_to_bytes32(42));
    let merkle = merkleize(data);

    assert_eq!(merkle, modified_cache);
}

#[test]
fn outer_builds() {
    let inner = Inner {
        a: 1,
        b: 2,
        c: 3,
        d: 4,
    };

    let outer = Outer {
        a: 0,
        b: inner.clone(),
        c: 5,
    };

    // Build the function output.
    let cache = outer.build_cache_bytes();

    // Generate reference data.
    let mut data = vec![];
    data.append(&mut int_to_bytes32(0));
    data.append(&mut inner.build_cache_bytes());
    data.append(&mut int_to_bytes32(5));
    let merkle = merkleize(data);

    assert_eq!(merkle, cache);
}

#[test]
fn partial_modification_u64_vec() {
    let n: u64 = 50;

    let original_vec: Vec<u64> = (0..n).collect();

    // Generate initial cache.
    let original_cache = original_vec.build_cache_bytes();

    // Modify the vec
    let mut modified_vec = original_vec.clone();
    modified_vec[n as usize - 1] = 42;

    // Perform a differential hash
    let mut cache_struct = TreeHashCache::from_bytes(original_cache.clone()).unwrap();
    modified_vec.cached_hash_tree_root(&original_vec, &mut cache_struct, 0);
    let modified_cache: Vec<u8> = cache_struct.into();

    // Generate reference data.
    let mut data = vec![];
    for i in &modified_vec {
        data.append(&mut int_to_bytes8(*i));
    }
    let data = sanitise_bytes(data);
    let expected = merkleize(data);

    assert_eq!(expected, modified_cache);
}

#[test]
fn large_vec_of_u64_builds() {
    let n: u64 = 50;

    let my_vec: Vec<u64> = (0..n).collect();

    // Generate function output.
    let cache = my_vec.build_cache_bytes();

    // Generate reference data.
    let mut data = vec![];
    for i in &my_vec {
        data.append(&mut int_to_bytes8(*i));
    }
    let data = sanitise_bytes(data);
    let expected = merkleize(data);

    assert_eq!(expected, cache);
}

#[test]
fn vec_of_u64_builds() {
    let data = join(vec![
        int_to_bytes8(1),
        int_to_bytes8(2),
        int_to_bytes8(3),
        int_to_bytes8(4),
        int_to_bytes8(5),
        vec![0; 32 - 8], // padding
    ]);

    let expected = merkleize(data);

    let my_vec = vec![1, 2, 3, 4, 5];

    let cache = my_vec.build_cache_bytes();

    assert_eq!(expected, cache);
}

#[test]
fn merkleize_odd() {
    let data = join(vec![
        int_to_bytes32(1),
        int_to_bytes32(2),
        int_to_bytes32(3),
        int_to_bytes32(4),
        int_to_bytes32(5),
    ]);

    let merkle = merkleize(sanitise_bytes(data));

    let expected_len = num_nodes(8) * BYTES_PER_CHUNK;

    assert_eq!(merkle.len(), expected_len);
}

fn generic_test(index: usize) {
    let inner = Inner {
        a: 1,
        b: 2,
        c: 3,
        d: 4,
    };

    let cache = inner.build_cache_bytes();

    let changed_inner = match index {
        0 => Inner {
            a: 42,
            ..inner.clone()
        },
        1 => Inner {
            b: 42,
            ..inner.clone()
        },
        2 => Inner {
            c: 42,
            ..inner.clone()
        },
        3 => Inner {
            d: 42,
            ..inner.clone()
        },
        _ => panic!("bad index"),
    };

    let mut cache_struct = TreeHashCache::from_bytes(cache.clone()).unwrap();

    changed_inner
        .cached_hash_tree_root(&inner, &mut cache_struct, 0)
        .unwrap();

    // assert_eq!(*cache_struct.hash_count, 3);

    let new_cache: Vec<u8> = cache_struct.into();

    let data1 = int_to_bytes32(1);
    let data2 = int_to_bytes32(2);
    let data3 = int_to_bytes32(3);
    let data4 = int_to_bytes32(4);

    let mut data = vec![data1, data2, data3, data4];

    data[index] = int_to_bytes32(42);

    let expected = merkleize(join(data));

    assert_eq!(expected, new_cache);
}

#[test]
fn cached_hash_on_inner() {
    generic_test(0);
    generic_test(1);
    generic_test(2);
    generic_test(3);
}

#[test]
fn inner_builds() {
    let data1 = int_to_bytes32(1);
    let data2 = int_to_bytes32(2);
    let data3 = int_to_bytes32(3);
    let data4 = int_to_bytes32(4);

    let data = join(vec![data1, data2, data3, data4]);
    let expected = merkleize(data);

    let inner = Inner {
        a: 1,
        b: 2,
        c: 3,
        d: 4,
    };

    let cache = inner.build_cache_bytes();

    assert_eq!(expected, cache);
}

#[test]
fn merkleize_4_leaves() {
    let data1 = hash(&int_to_bytes32(1));
    let data2 = hash(&int_to_bytes32(2));
    let data3 = hash(&int_to_bytes32(3));
    let data4 = hash(&int_to_bytes32(4));

    let data = join(vec![
        data1.clone(),
        data2.clone(),
        data3.clone(),
        data4.clone(),
    ]);

    let cache = merkleize(data);

    let hash_12 = {
        let mut joined = vec![];
        joined.append(&mut data1.clone());
        joined.append(&mut data2.clone());
        hash(&joined)
    };
    let hash_34 = {
        let mut joined = vec![];
        joined.append(&mut data3.clone());
        joined.append(&mut data4.clone());
        hash(&joined)
    };
    let hash_hash12_hash_34 = {
        let mut joined = vec![];
        joined.append(&mut hash_12.clone());
        joined.append(&mut hash_34.clone());
        hash(&joined)
    };

    for (i, chunk) in cache.chunks(HASHSIZE).enumerate().rev() {
        let expected = match i {
            0 => hash_hash12_hash_34.clone(),
            1 => hash_12.clone(),
            2 => hash_34.clone(),
            3 => data1.clone(),
            4 => data2.clone(),
            5 => data3.clone(),
            6 => data4.clone(),
            _ => vec![],
        };

        assert_eq!(chunk, &expected[..], "failed at {}", i);
    }
}
<|endoftext|>


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

#[macro_use]
extern crate gfx;
extern crate gfx_window_glutin;
extern crate glutin;

pub use gfx::format::{Srgb8, Depth};

gfx_vertex_struct!( Vertex {
    pos: [f32; 2] = "a_Pos",
    uv: [f32; 2] = "a_Uv",
});

impl Vertex {
    fn new(p: [f32; 2], u: [f32; 2]) -> Vertex {
        Vertex {
            pos: p,
            uv: u,
        }
    }
}

gfx_pipeline!(pipe {
    vbuf: gfx::VertexBuffer<Vertex> = (),
    tex: gfx::TextureSampler<[f32; 3]> = "t_Tex",
    out: gfx::RenderTarget<Srgb8> = "o_Color",
});

// Larger red dots
const L0_DATA: [u8; 16] = [
    0x00, 0x00, 0x00, 0x00,
    0x00, 0xc0, 0xc0, 0x00,
    0x00, 0xc0, 0xc0, 0x00,
    0x00, 0x00, 0x00, 0x00,
];

// Uniform green
const L1_DATA: [u8; 4] = [
    0x10, 0x18,
    0x18, 0x18,
];

// Uniform blue
const L2_DATA: [u8; 1] = [ 0x02 ];

fn make_texture<R, F>(factory: &mut F) -> gfx::handle::ShaderResourceView<R, [f32; 3]>
        where R: gfx::Resources,
              F: gfx::Factory<R>
{
    let kind = gfx::tex::Kind::D2(4, 4, gfx::tex::AaMode::Single);
    let tex = factory.create_texture(kind, 3, gfx::SHADER_RESOURCE,
        Some(gfx::format::ChannelType::Unorm)).unwrap();

    type Rgb332 = (gfx::format::R3_G3_B2, gfx::format::Unorm);

    factory.update_texture::<Rgb332>(&tex, &tex.get_info().to_image_info(0),
        gfx::cast_slice(&L0_DATA), None).unwrap();
    factory.update_texture::<Rgb332>(&tex, &tex.get_info().to_image_info(1),
        gfx::cast_slice(&L1_DATA), None).unwrap();
    factory.update_texture::<Rgb332>(&tex, &tex.get_info().to_image_info(2),
        gfx::cast_slice(&L2_DATA), None).unwrap();

    factory.view_texture_as_shader_resource::<Rgb332>(
        &tex, (0, 2), gfx::format::Swizzle::new()).unwrap()
}

pub fn main() {
    use gfx::traits::{Device, Factory, FactoryExt};

    let builder = glutin::WindowBuilder::new()
        .with_title("Mipmap example".to_string())
        .with_dimensions(800, 600);
    let (window, mut device, mut factory, main_color, _) =
        gfx_window_glutin::init::<Srgb8, Depth>(builder);
    let mut encoder = factory.create_encoder();

    let pso = factory.create_pipeline_simple(
        include_bytes!("shader/120.glslv"),
        include_bytes!("shader/120.glslf"),
        gfx::state::CullFace::Nothing,
        pipe::new()
        ).unwrap();

    let vertex_data = [
        Vertex::new([ 0.0,  0.0], [ 0.0,  0.0]),
        Vertex::new([ 1.0,  0.0], [50.0,  0.0]),
        Vertex::new([ 1.0,  1.1], [50.0, 50.0]),

        Vertex::new([ 0.0,  0.0], [  0.0,   0.0]),
        Vertex::new([-1.0,  0.0], [800.0,   0.0]),
        Vertex::new([-1.0, -1.0], [800.0, 800.0]),
    ];
    let (vbuf, slice) = factory.create_vertex_buffer(&vertex_data);

    let texture_view = make_texture(&mut factory);

    let sampler = factory.create_sampler(gfx::tex::SamplerInfo::new(
        gfx::tex::FilterMethod::Trilinear,
        gfx::tex::WrapMode::Tile,
    ));

    let data = pipe::Data {
        vbuf: vbuf,
        tex: (texture_view, sampler),
        out: main_color,
    };

    'main: loop {
        // quit when Esc is pressed.
        for event in window.poll_events() {
            match event {
                glutin::Event::KeyboardInput(_, _, Some(glutin::VirtualKeyCode::Escape)) |
                glutin::Event::Closed => break 'main,
                _ => {},
            }
        }

        encoder.reset();
        encoder.clear(&data.out, [0.1, 0.2, 0.3]);
        encoder.draw(&slice, &pso, &data);

        device.submit(encoder.as_buffer());
        window.swap_buffers().unwrap();
        device.cleanup();
    }
}
<commit_msg>Auto merge of #893 - vickenty:mipmap-rgb8, r=sectopod<commit_after>// Copyright 2015 The Gfx-rs Developers.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#[macro_use]
extern crate gfx;
extern crate gfx_window_glutin;
extern crate glutin;

pub use gfx::format::{Rgba8, Depth};

gfx_vertex_struct!( Vertex {
    pos: [f32; 2] = "a_Pos",
    uv: [f32; 2] = "a_Uv",
});

impl Vertex {
    fn new(p: [f32; 2], u: [f32; 2]) -> Vertex {
        Vertex {
            pos: p,
            uv: u,
        }
    }
}

gfx_pipeline!(pipe {
    vbuf: gfx::VertexBuffer<Vertex> = (),
    tex: gfx::TextureSampler<[f32; 4]> = "t_Tex",
    out: gfx::RenderTarget<Rgba8> = "o_Color",
});

// Larger red dots
const L0_DATA: [[u8; 4]; 16] = [
    [ 0x00, 0x00, 0x00, 0x00 ], [ 0x00, 0x00, 0x00, 0x00 ], [ 0x00, 0x00, 0x00, 0x00 ], [ 0x00, 0x00, 0x00, 0x00 ],
    [ 0x00, 0x00, 0x00, 0x00 ], [ 0xc0, 0x00, 0x00, 0x00 ], [ 0xc0, 0x00, 0x00, 0x00 ], [ 0x00, 0x00, 0x00, 0x00 ],
    [ 0x00, 0x00, 0x00, 0x00 ], [ 0xc0, 0x00, 0x00, 0x00 ], [ 0xc0, 0x00, 0x00, 0x00 ], [ 0x00, 0x00, 0x00, 0x00 ],
    [ 0x00, 0x00, 0x00, 0x00 ], [ 0x00, 0x00, 0x00, 0x00 ], [ 0x00, 0x00, 0x00, 0x00 ], [ 0x00, 0x00, 0x00, 0x00 ],
];

// Uniform green
const L1_DATA: [[u8; 4]; 4] = [
    [ 0x00, 0xc0, 0x00, 0x00 ], [ 0x00, 0xc0, 0x00, 0x00 ],
    [ 0x00, 0xc0, 0x00, 0x00 ], [ 0x00, 0xc0, 0x00, 0x00 ],
];

// Uniform blue
const L2_DATA: [[u8; 4]; 1] = [ [ 0x00, 0x00, 0xc0, 0x00 ] ];

fn make_texture<R, F>(factory: &mut F) -> gfx::handle::ShaderResourceView<R, [f32; 4]>
        where R: gfx::Resources,
              F: gfx::Factory<R>
{
    let kind = gfx::tex::Kind::D2(4, 4, gfx::tex::AaMode::Single);
    let tex = factory.create_texture(kind, 3, gfx::SHADER_RESOURCE,
        Some(gfx::format::ChannelType::Unorm)).unwrap();

    factory.update_texture::<Rgba8>(&tex, &tex.get_info().to_image_info(0),
        gfx::cast_slice(&L0_DATA), None).unwrap();
    factory.update_texture::<Rgba8>(&tex, &tex.get_info().to_image_info(1),
        gfx::cast_slice(&L1_DATA), None).unwrap();
    factory.update_texture::<Rgba8>(&tex, &tex.get_info().to_image_info(2),
        gfx::cast_slice(&L2_DATA), None).unwrap();

    factory.view_texture_as_shader_resource::<Rgba8>(
        &tex, (0, 2), gfx::format::Swizzle::new()).unwrap()
}

pub fn main() {
    use gfx::traits::{Device, Factory, FactoryExt};

    let builder = glutin::WindowBuilder::new()
        .with_title("Mipmap example".to_string())
        .with_dimensions(800, 600);
    let (window, mut device, mut factory, main_color, _) =
        gfx_window_glutin::init::<Rgba8, Depth>(builder);
    let mut encoder = factory.create_encoder();

    let pso = factory.create_pipeline_simple(
        include_bytes!("shader/120.glslv"),
        include_bytes!("shader/120.glslf"),
        gfx::state::CullFace::Nothing,
        pipe::new()
        ).unwrap();

    let vertex_data = [
        Vertex::new([ 0.0,  0.0], [ 0.0,  0.0]),
        Vertex::new([ 1.0,  0.0], [50.0,  0.0]),
        Vertex::new([ 1.0,  1.1], [50.0, 50.0]),

        Vertex::new([ 0.0,  0.0], [  0.0,   0.0]),
        Vertex::new([-1.0,  0.0], [800.0,   0.0]),
        Vertex::new([-1.0, -1.0], [800.0, 800.0]),
    ];
    let (vbuf, slice) = factory.create_vertex_buffer(&vertex_data);

    let texture_view = make_texture(&mut factory);

    let sampler = factory.create_sampler(gfx::tex::SamplerInfo::new(
        gfx::tex::FilterMethod::Trilinear,
        gfx::tex::WrapMode::Tile,
    ));

    let data = pipe::Data {
        vbuf: vbuf,
        tex: (texture_view, sampler),
        out: main_color,
    };

    'main: loop {
        // quit when Esc is pressed.
        for event in window.poll_events() {
            match event {
                glutin::Event::KeyboardInput(_, _, Some(glutin::VirtualKeyCode::Escape)) |
                glutin::Event::Closed => break 'main,
                _ => {},
            }
        }

        encoder.reset();
        encoder.clear(&data.out, [0.1, 0.2, 0.3, 1.0]);
        encoder.draw(&slice, &pso, &data);

        device.submit(encoder.as_buffer());
        window.swap_buffers().unwrap();
        device.cleanup();
    }
}
<|endoftext|>