<commit_before>use std::collections::HashSet;

use color::Color;
use util::{Coords, lift_tuple};

pub struct Program {
    size: Coords,
    image: Vec<Color>,
}

impl Program {
    pub fn new(size: Coords, image: Vec<Color>) -> Self {
        Program {
            size: size,
            image: image,
        }
    }

    pub fn get(&self, coords: Coords) -> Option<Color> {
        self.coords_to_index(coords).map(|index| self.image[index])
    }

    /// Find the coordinates of a contiguous area of codels of the same color,
    /// starting from a coordinate.
    ///
    /// TODO: This implementation is recursive, probably not the most optimal.
    /// Fix it if necessary. (Or if you've had enough of the terrible code.)
    /// See also: https://en.wikipedia.org/wiki/Flood_fill
    pub fn color_block(&self, start_coords: Coords) -> HashSet<Coords> {
        let mut result = HashSet::new();

        let start_color = self.get(start_coords).unwrap();
        // TODO: decide if this needs actual error handling

        self.mark_blocks(start_coords, start_color, &mut result);

        result
    }

    /// The recursive part of `color_block`.
    fn mark_blocks(&self, coords: Coords, start_color: Color, marked: &mut HashSet<Coords>) {
        if marked.contains(&coords) {
            return;
        }

        if let Some(color) = self.get(coords) {
            if color != start_color {
                return;
            }

            marked.insert(coords);

            for neighbor in self.neighbors(coords).iter().filter_map(|&x| x) {
                self.mark_blocks(neighbor, start_color, marked);
            }
        }
    }

    fn neighbors(&self, coords: Coords) -> [Option<Coords>; 4] {
        let (x, y) = coords;

        let right = lift_tuple((Some(x + 1), Some(y)));
        let left = lift_tuple((x.checked_sub(1), Some(y)));
        let above = lift_tuple((Some(x), Some(y + 1)));
        let below = lift_tuple((Some(x), y.checked_sub(1)));

        // I would have written this with iterators, but AFAIK it's not possible
        // to `.collect()` into a fixed size array.
        [right.and_then(|c| self.check_coords(c)),
         left.and_then(|c| self.check_coords(c)),
         above.and_then(|c| self.check_coords(c)),
         below.and_then(|c| self.check_coords(c))]
    }

    fn coords_to_index(&self, coords: Coords) -> Option<usize> {
        let (x, y) = coords;
        let (width, height) = self.size;

        self.check_coords(coords).and(Some(width * y + x % width))
    }

    fn index_to_coords(&self, index: usize) -> Option<Coords> {
        let (width, height) = self.size;

        if index >= width * height {
            return None;
        }

        let x = index % width;
        let y = index / width;

        Some((x, y))
    }

    fn check_coords(&self, coords: Coords) -> Option<Coords> {
        let (x, y) = coords;
        let (width, height) = self.size;

        if x >= width || y >= height {
            None
        } else {
            Some(coords)
        }
    }
}

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

    #[test]
    fn test_program_coords_to_index() {
        let program = Program::new((5, 5), vec![]);
        assert_eq!(program.coords_to_index((0, 0)), Some(0));
        assert_eq!(program.coords_to_index((2, 0)), Some(2));
        assert_eq!(program.coords_to_index((0, 2)), Some(10));
        assert_eq!(program.coords_to_index((2, 2)), Some(12));
        assert_eq!(program.coords_to_index((4, 4)), Some(24));
        assert_eq!(program.coords_to_index((5, 5)), None);
    }

    #[test]
    fn test_program_index_to_coords() {
        let program = Program::new((5, 5), vec![]);
        assert_eq!(program.index_to_coords(0), Some((0, 0)));
        assert_eq!(program.index_to_coords(2), Some((2, 0)));
        assert_eq!(program.index_to_coords(10), Some((0, 2)));
        assert_eq!(program.index_to_coords(12), Some((2, 2)));
        assert_eq!(program.index_to_coords(24), Some((4, 4)));
        assert_eq!(program.index_to_coords(25), None);
    }

    #[test]
    fn test_program_check_coords() {
        let program = Program::new((5, 5), vec![]);
        assert!(program.check_coords((0, 0)).is_some());
        assert!(program.check_coords((4, 0)).is_some());
        assert!(program.check_coords((0, 4)).is_some());
        assert!(program.check_coords((4, 4)).is_some());
        assert!(program.check_coords((5, 4)).is_none());
        assert!(program.check_coords((4, 5)).is_none());
        assert!(program.check_coords((5, 5)).is_none());
    }

    #[test]
    #[cfg_attr(rustfmt, rustfmt_skip)]
    fn test_program_color_block() {
        use super::super::Color::White as W;
        use super::super::Color::Black as B;

        let mk_prog = |p| Program::new((5, 5), p);

        let program = mk_prog(vec![
            W, W, W, W, W,
            W, W, W, W, W,
            W, W, B, W, W,
            W, W, W, W, W,
            W, W, W, W, W,
        ]);
        let block = vec![(2, 2)].into_iter().collect();
        assert_eq!(program.color_block((2, 2)), block);

        let program = mk_prog( vec![
            B, W, W, W, W,
            W, W, W, W, W,
            W, W, W, W, W,
            W, W, W, W, W,
            W, W, W, W, W,
        ]);
        let block = vec![(0, 0)].into_iter().collect();
        assert_eq!(program.color_block((0, 0)), block);

        let program = mk_prog(vec![
            W, W, W, W, W,
            W, B, B, B, W,
            W, B, B, B, W,
            W, B, B, B, W,
            W, W, W, W, W,
        ]);
        let block = vec![
            (1, 1), (2, 1), (3, 1),
            (1, 2), (2, 2), (3, 2),
            (1, 3), (2, 3), (3, 3),
        ].into_iter().collect();
        assert_eq!(program.color_block((2, 2)), block);

        let program = mk_prog(vec![
            W, W, W, W, W,
            W, B, B, B, W,
            W, B, W, B, W,
            W, B, W, B, W,
            W, W, W, W, W,
        ]);
        let block = vec![
            (1, 1), (2, 1), (3, 1),
            (1, 2),         (3, 2),
            (1, 3),         (3, 3),
        ].into_iter().collect();
        assert_eq!(program.color_block((1, 1)), block);

        let program = mk_prog(vec![
            W, W, W, W, W,
            W, B, W, B, W,
            W, B, W, B, W,
            W, B, W, B, W,
            W, W, W, W, W,
        ]);
        let block = vec![
            (1, 1),
            (1, 2),
            (1, 3),
        ].into_iter().collect();
        assert_eq!(program.color_block((1, 1)), block);
    }
}
<commit_msg>Derive Debug for Program<commit_after>use std::collections::HashSet;

use color::Color;
use util::{Coords, lift_tuple};

#[derive(Debug)]
pub struct Program {
    size: Coords,
    image: Vec<Color>,
}

impl Program {
    pub fn new(size: Coords, image: Vec<Color>) -> Self {
        Program {
            size: size,
            image: image,
        }
    }

    pub fn get(&self, coords: Coords) -> Option<Color> {
        self.coords_to_index(coords).map(|index| self.image[index])
    }

    /// Find the coordinates of a contiguous area of codels of the same color,
    /// starting from a coordinate.
    ///
    /// TODO: This implementation is recursive, probably not the most optimal.
    /// Fix it if necessary. (Or if you've had enough of the terrible code.)
    /// See also: https://en.wikipedia.org/wiki/Flood_fill
    pub fn color_block(&self, start_coords: Coords) -> HashSet<Coords> {
        let mut result = HashSet::new();

        let start_color = self.get(start_coords).unwrap();
        // TODO: decide if this needs actual error handling

        self.mark_blocks(start_coords, start_color, &mut result);

        result
    }

    /// The recursive part of `color_block`.
    fn mark_blocks(&self, coords: Coords, start_color: Color, marked: &mut HashSet<Coords>) {
        if marked.contains(&coords) {
            return;
        }

        if let Some(color) = self.get(coords) {
            if color != start_color {
                return;
            }

            marked.insert(coords);

            for neighbor in self.neighbors(coords).iter().filter_map(|&x| x) {
                self.mark_blocks(neighbor, start_color, marked);
            }
        }
    }

    fn neighbors(&self, coords: Coords) -> [Option<Coords>; 4] {
        let (x, y) = coords;

        let right = lift_tuple((Some(x + 1), Some(y)));
        let left = lift_tuple((x.checked_sub(1), Some(y)));
        let above = lift_tuple((Some(x), Some(y + 1)));
        let below = lift_tuple((Some(x), y.checked_sub(1)));

        // I would have written this with iterators, but AFAIK it's not possible
        // to `.collect()` into a fixed size array.
        [right.and_then(|c| self.check_coords(c)),
         left.and_then(|c| self.check_coords(c)),
         above.and_then(|c| self.check_coords(c)),
         below.and_then(|c| self.check_coords(c))]
    }

    fn coords_to_index(&self, coords: Coords) -> Option<usize> {
        let (x, y) = coords;
        let (width, height) = self.size;

        self.check_coords(coords).and(Some(width * y + x % width))
    }

    fn index_to_coords(&self, index: usize) -> Option<Coords> {
        let (width, height) = self.size;

        if index >= width * height {
            return None;
        }

        let x = index % width;
        let y = index / width;

        Some((x, y))
    }

    fn check_coords(&self, coords: Coords) -> Option<Coords> {
        let (x, y) = coords;
        let (width, height) = self.size;

        if x >= width || y >= height {
            None
        } else {
            Some(coords)
        }
    }
}

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

    #[test]
    fn test_program_coords_to_index() {
        let program = Program::new((5, 5), vec![]);
        assert_eq!(program.coords_to_index((0, 0)), Some(0));
        assert_eq!(program.coords_to_index((2, 0)), Some(2));
        assert_eq!(program.coords_to_index((0, 2)), Some(10));
        assert_eq!(program.coords_to_index((2, 2)), Some(12));
        assert_eq!(program.coords_to_index((4, 4)), Some(24));
        assert_eq!(program.coords_to_index((5, 5)), None);
    }

    #[test]
    fn test_program_index_to_coords() {
        let program = Program::new((5, 5), vec![]);
        assert_eq!(program.index_to_coords(0), Some((0, 0)));
        assert_eq!(program.index_to_coords(2), Some((2, 0)));
        assert_eq!(program.index_to_coords(10), Some((0, 2)));
        assert_eq!(program.index_to_coords(12), Some((2, 2)));
        assert_eq!(program.index_to_coords(24), Some((4, 4)));
        assert_eq!(program.index_to_coords(25), None);
    }

    #[test]
    fn test_program_check_coords() {
        let program = Program::new((5, 5), vec![]);
        assert!(program.check_coords((0, 0)).is_some());
        assert!(program.check_coords((4, 0)).is_some());
        assert!(program.check_coords((0, 4)).is_some());
        assert!(program.check_coords((4, 4)).is_some());
        assert!(program.check_coords((5, 4)).is_none());
        assert!(program.check_coords((4, 5)).is_none());
        assert!(program.check_coords((5, 5)).is_none());
    }

    #[test]
    #[cfg_attr(rustfmt, rustfmt_skip)]
    fn test_program_color_block() {
        use super::super::Color::White as W;
        use super::super::Color::Black as B;

        let mk_prog = |p| Program::new((5, 5), p);

        let program = mk_prog(vec![
            W, W, W, W, W,
            W, W, W, W, W,
            W, W, B, W, W,
            W, W, W, W, W,
            W, W, W, W, W,
        ]);
        let block = vec![(2, 2)].into_iter().collect();
        assert_eq!(program.color_block((2, 2)), block);

        let program = mk_prog( vec![
            B, W, W, W, W,
            W, W, W, W, W,
            W, W, W, W, W,
            W, W, W, W, W,
            W, W, W, W, W,
        ]);
        let block = vec![(0, 0)].into_iter().collect();
        assert_eq!(program.color_block((0, 0)), block);

        let program = mk_prog(vec![
            W, W, W, W, W,
            W, B, B, B, W,
            W, B, B, B, W,
            W, B, B, B, W,
            W, W, W, W, W,
        ]);
        let block = vec![
            (1, 1), (2, 1), (3, 1),
            (1, 2), (2, 2), (3, 2),
            (1, 3), (2, 3), (3, 3),
        ].into_iter().collect();
        assert_eq!(program.color_block((2, 2)), block);

        let program = mk_prog(vec![
            W, W, W, W, W,
            W, B, B, B, W,
            W, B, W, B, W,
            W, B, W, B, W,
            W, W, W, W, W,
        ]);
        let block = vec![
            (1, 1), (2, 1), (3, 1),
            (1, 2),         (3, 2),
            (1, 3),         (3, 3),
        ].into_iter().collect();
        assert_eq!(program.color_block((1, 1)), block);

        let program = mk_prog(vec![
            W, W, W, W, W,
            W, B, W, B, W,
            W, B, W, B, W,
            W, B, W, B, W,
            W, W, W, W, W,
        ]);
        let block = vec![
            (1, 1),
            (1, 2),
            (1, 3),
        ].into_iter().collect();
        assert_eq!(program.color_block((1, 1)), block);
    }
}
<|endoftext|>


<commit_before>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::{
    backup_types::epoch_ending::manifest::EpochEndingBackup,
    metrics::{
        restore::{EPOCH_ENDING_EPOCH, EPOCH_ENDING_VERSION},
        verify::{VERIFY_EPOCH_ENDING_EPOCH, VERIFY_EPOCH_ENDING_VERSION},
    },
    storage::{BackupStorage, FileHandle, FileHandleRef},
    utils::{
        read_record_bytes::ReadRecordBytes, storage_ext::BackupStorageExt, stream::StreamX,
        GlobalRestoreOptions, RestoreRunMode,
    },
};
use anyhow::{anyhow, ensure, Result};
use aptos_logger::prelude::*;
use aptos_types::{
    epoch_change::Verifier,
    ledger_info::{LedgerInfo, LedgerInfoWithSignatures},
    transaction::Version,
    waypoint::Waypoint,
};
use clap::Parser;
use futures::StreamExt;
use std::{collections::HashMap, sync::Arc, time::Instant};

#[derive(Parser)]
pub struct EpochEndingRestoreOpt {
    #[clap(long = "epoch-ending-manifest")]
    pub manifest_handle: FileHandle,
}

pub struct EpochEndingRestoreController {
    storage: Arc<dyn BackupStorage>,
    run_mode: Arc<RestoreRunMode>,
    manifest_handle: FileHandle,
    target_version: Version,
    trusted_waypoints: Arc<HashMap<Version, Waypoint>>,
}

impl EpochEndingRestoreController {
    pub fn new(
        opt: EpochEndingRestoreOpt,
        global_opt: GlobalRestoreOptions,
        storage: Arc<dyn BackupStorage>,
    ) -> Self {
        Self {
            storage,
            run_mode: global_opt.run_mode,
            manifest_handle: opt.manifest_handle,
            target_version: global_opt.target_version,
            trusted_waypoints: global_opt.trusted_waypoints,
        }
    }

    pub async fn preheat(self) -> PreheatedEpochEndingRestore {
        PreheatedEpochEndingRestore {
            preheat_result: self.preheat_impl().await,
            controller: self,
        }
    }

    pub async fn run(
        self,
        previous_epoch_ending_ledger_info: Option<&LedgerInfo>,
    ) -> Result<Vec<LedgerInfo>> {
        self.preheat()
            .await
            .run(previous_epoch_ending_ledger_info)
            .await
    }
}

impl EpochEndingRestoreController {
    fn name(&self) -> String {
        format!("epoch ending {}", self.run_mode.name())
    }

    async fn preheat_impl(&self) -> Result<EpochEndingRestorePreheatData> {
        let manifest: EpochEndingBackup =
            self.storage.load_json_file(&self.manifest_handle).await?;
        manifest.verify()?;

        let mut next_epoch = manifest.first_epoch;
        let mut waypoint_iter = manifest.waypoints.iter();

        let mut previous_li: Option<&LedgerInfoWithSignatures> = None;
        let mut ledger_infos = Vec::new();

        let mut past_target = false;
        for chunk in &manifest.chunks {
            if past_target {
                break;
            }

            let lis = self.read_chunk(&chunk.ledger_infos).await?;
            ensure!(
                chunk.first_epoch + lis.len() as u64 == chunk.last_epoch + 1,
                "Number of items in chunks doesn't match that in manifest. \
                first_epoch: {}, last_epoch: {}, items in chunk: {}",
                chunk.first_epoch,
                chunk.last_epoch,
                lis.len(),
            );

            for li in lis {
                if li.ledger_info().version() > self.target_version {
                    past_target = true;
                    break;
                }

                ensure!(
                    li.ledger_info().epoch() == next_epoch,
                    "LedgerInfo epoch not expected. Expected: {}, actual: {}.",
                    li.ledger_info().epoch(),
                    next_epoch,
                );
                let wp_manifest = waypoint_iter.next().ok_or_else(|| {
                    anyhow!("More LedgerInfo's found than waypoints in manifest.")
                })?;
                let wp_li = Waypoint::new_epoch_boundary(li.ledger_info())?;
                ensure!(
                    *wp_manifest == wp_li,
                    "Waypoints don't match. In manifest: {}, In chunk: {}",
                    wp_manifest,
                    wp_li,
                );
                if let Some(wp_trusted) = self.trusted_waypoints.get(&wp_li.version()) {
                    ensure!(
                        *wp_trusted == wp_li,
                        "Waypoints don't match. In backup: {}, trusted: {}",
                        wp_li,
                        wp_trusted,
                    );
                } else if let Some(pre_li) = previous_li {
                    pre_li
                        .ledger_info()
                        .next_epoch_state()
                        .ok_or_else(|| {
                            anyhow!(
                                "Next epoch state not found from LI at epoch {}.",
                                pre_li.ledger_info().epoch()
                            )
                        })?
                        .verify(&li)?;
                }
                ledger_infos.push(li);
                previous_li = ledger_infos.last();
                next_epoch += 1;
            }
        }

        Ok(EpochEndingRestorePreheatData {
            manifest,
            ledger_infos,
        })
    }

    async fn read_chunk(
        &self,
        file_handle: &FileHandleRef,
    ) -> Result<Vec<LedgerInfoWithSignatures>> {
        let mut file = self.storage.open_for_read(file_handle).await?;
        let mut chunk = vec![];

        while let Some(record_bytes) = file.read_record_bytes().await? {
            chunk.push(bcs::from_bytes(&record_bytes)?);
        }

        Ok(chunk)
    }
}

struct EpochEndingRestorePreheatData {
    manifest: EpochEndingBackup,
    ledger_infos: Vec<LedgerInfoWithSignatures>,
}

pub struct PreheatedEpochEndingRestore {
    controller: EpochEndingRestoreController,
    preheat_result: Result<EpochEndingRestorePreheatData>,
}

impl PreheatedEpochEndingRestore {
    pub async fn run(
        self,
        previous_epoch_ending_ledger_info: Option<&LedgerInfo>,
    ) -> Result<Vec<LedgerInfo>> {
        let name = self.controller.name();
        info!(
            "{} started. Manifest: {}",
            name, self.controller.manifest_handle
        );
        let res = self
            .run_impl(previous_epoch_ending_ledger_info)
            .await
            .map_err(|e| anyhow!("{} failed: {}", name, e))?;
        info!("{} succeeded.", name);
        Ok(res)
    }
}

impl PreheatedEpochEndingRestore {
    async fn run_impl(
        self,
        previous_epoch_ending_ledger_info: Option<&LedgerInfo>,
    ) -> Result<Vec<LedgerInfo>> {
        let preheat_data = self
            .preheat_result
            .map_err(|e| anyhow!("Preheat failed: {}", e))?;

        let first_li = preheat_data
            .ledger_infos
            .first()
            .expect("Epoch ending backup can't be empty.");

        if let Some(li) = previous_epoch_ending_ledger_info {
            ensure!(
                li.next_block_epoch() == preheat_data.manifest.first_epoch,
                "Previous epoch ending LedgerInfo is not the one expected. \
                My first epoch: {}, previous LedgerInfo next_block_epoch: {}",
                preheat_data.manifest.first_epoch,
                li.next_block_epoch(),
            );
            // Waypoint has been verified in preheat if it's trusted, otherwise try to check
            // the signatures.
            if self
                .controller
                .trusted_waypoints
                .get(&first_li.ledger_info().version())
                .is_none()
            {
                li.next_epoch_state()
                    .ok_or_else(|| {
                        anyhow!("Previous epoch ending LedgerInfo doesn't end an epoch")
                    })?
                    .verify(first_li)?;
            }
        }

        let last_li = preheat_data
            .ledger_infos
            .last()
            .expect("Verified not empty.")
            .ledger_info();
        match self.controller.run_mode.as_ref() {
            RestoreRunMode::Restore { restore_handler } => {
                restore_handler.save_ledger_infos(&preheat_data.ledger_infos)?;

                EPOCH_ENDING_EPOCH.set(last_li.epoch() as i64);
                EPOCH_ENDING_VERSION.set(last_li.version() as i64);
            }
            RestoreRunMode::Verify => {
                VERIFY_EPOCH_ENDING_EPOCH.set(last_li.epoch() as i64);
                VERIFY_EPOCH_ENDING_VERSION.set(last_li.version() as i64);
            }
        };

        Ok(preheat_data
            .ledger_infos
            .into_iter()
            .map(|x| x.ledger_info().clone())
            .collect())
    }
}

/// Represents a history of epoch changes since epoch 0.
#[derive(Clone)]
pub struct EpochHistory {
    pub epoch_endings: Vec<LedgerInfo>,
}

impl EpochHistory {
    pub fn verify_ledger_info(&self, li_with_sigs: &LedgerInfoWithSignatures) -> Result<()> {
        let epoch = li_with_sigs.ledger_info().epoch();
        ensure!(!self.epoch_endings.is_empty(), "Empty epoch history.",);
        if epoch > self.epoch_endings.len() as u64 {
            // TODO(aldenhu): fix this from upper level
            warn!(
                epoch = epoch,
                epoch_history_until = self.epoch_endings.len(),
                "Epoch is too new and can't be verified. Previous chunks are verified and node \
                won't be able to start if this data is malicious."
            );
            return Ok(());
        }
        if epoch == 0 {
            ensure!(
                li_with_sigs.ledger_info() == &self.epoch_endings[0],
                "Genesis epoch LedgerInfo info doesn't match.",
            );
            Ok(())
        } else {
            self.epoch_endings[epoch as usize - 1]
                .next_epoch_state()
                .ok_or_else(|| anyhow!("Shouldn't contain non- epoch bumping LIs."))?
                .verify(li_with_sigs)
        }
    }
}

pub struct EpochHistoryRestoreController {
    storage: Arc<dyn BackupStorage>,
    manifest_handles: Vec<FileHandle>,
    global_opt: GlobalRestoreOptions,
}

impl EpochHistoryRestoreController {
    pub fn new(
        manifest_handles: Vec<FileHandle>,
        global_opt: GlobalRestoreOptions,
        storage: Arc<dyn BackupStorage>,
    ) -> Self {
        Self {
            storage,
            manifest_handles,
            global_opt,
        }
    }

    pub async fn run(self) -> Result<EpochHistory> {
        let name = self.name();
        info!("{} started.", name,);
        let res = self
            .run_impl()
            .await
            .map_err(|e| anyhow!("{} failed: {}", name, e))?;
        info!("{} succeeded.", name);
        Ok(res)
    }
}

impl EpochHistoryRestoreController {
    fn name(&self) -> String {
        format!("epoch history {}", self.global_opt.run_mode.name())
    }

    async fn run_impl(self) -> Result<EpochHistory> {
        let timer = Instant::now();
        if self.manifest_handles.is_empty() {
            return Ok(EpochHistory {
                epoch_endings: Vec::new(),
            });
        }

        let futs_iter = self.manifest_handles.iter().map(|hdl| {
            EpochEndingRestoreController::new(
                EpochEndingRestoreOpt {
                    manifest_handle: hdl.clone(),
                },
                self.global_opt.clone(),
                self.storage.clone(),
            )
            .preheat()
        });
        let mut futs_stream = futures::stream::iter(futs_iter).buffered_x(
            self.global_opt.concurrent_downloads * 2, /* buffer size */
            self.global_opt.concurrent_downloads,     /* concurrency */
        );

        let mut next_epoch = 0u64;
        let mut previous_li = None;
        let mut epoch_endings = Vec::new();

        while let Some(preheated_restore) = futs_stream.next().await {
            let manifest_handle = preheated_restore.controller.manifest_handle.clone();
            let lis = preheated_restore.run(previous_li).await?;
            ensure!(
                !lis.is_empty(),
                "No epochs restored from {}",
                manifest_handle,
            );
            for li in &lis {
                ensure!(
                    li.epoch() == next_epoch,
                    "Restored LedgerInfo has epoch {}, expecting {}.",
                    li.epoch(),
                    next_epoch,
                );
                ensure!(
                    li.ends_epoch(),
                    "LedgerInfo is not one at an epoch ending. epoch: {}",
                    li.epoch(),
                );
                next_epoch += 1;
            }

            epoch_endings.extend(lis);
            previous_li = epoch_endings.last();
        }

        info!(
            "Epoch history recovered in {:.2} seconds",
            timer.elapsed().as_secs_f64()
        );
        Ok(EpochHistory { epoch_endings })
    }
}
<commit_msg>[backup] still consider trusted-waypoints after EpochHistory is constructed<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::{
    backup_types::epoch_ending::manifest::EpochEndingBackup,
    metrics::{
        restore::{EPOCH_ENDING_EPOCH, EPOCH_ENDING_VERSION},
        verify::{VERIFY_EPOCH_ENDING_EPOCH, VERIFY_EPOCH_ENDING_VERSION},
    },
    storage::{BackupStorage, FileHandle, FileHandleRef},
    utils::{
        read_record_bytes::ReadRecordBytes, storage_ext::BackupStorageExt, stream::StreamX,
        GlobalRestoreOptions, RestoreRunMode,
    },
};
use anyhow::{anyhow, ensure, Result};
use aptos_logger::prelude::*;
use aptos_types::{
    epoch_change::Verifier,
    ledger_info::{LedgerInfo, LedgerInfoWithSignatures},
    transaction::Version,
    waypoint::Waypoint,
};
use clap::Parser;
use futures::StreamExt;
use std::{collections::HashMap, sync::Arc, time::Instant};

#[derive(Parser)]
pub struct EpochEndingRestoreOpt {
    #[clap(long = "epoch-ending-manifest")]
    pub manifest_handle: FileHandle,
}

pub struct EpochEndingRestoreController {
    storage: Arc<dyn BackupStorage>,
    run_mode: Arc<RestoreRunMode>,
    manifest_handle: FileHandle,
    target_version: Version,
    trusted_waypoints: Arc<HashMap<Version, Waypoint>>,
}

impl EpochEndingRestoreController {
    pub fn new(
        opt: EpochEndingRestoreOpt,
        global_opt: GlobalRestoreOptions,
        storage: Arc<dyn BackupStorage>,
    ) -> Self {
        Self {
            storage,
            run_mode: global_opt.run_mode,
            manifest_handle: opt.manifest_handle,
            target_version: global_opt.target_version,
            trusted_waypoints: global_opt.trusted_waypoints,
        }
    }

    pub async fn preheat(self) -> PreheatedEpochEndingRestore {
        PreheatedEpochEndingRestore {
            preheat_result: self.preheat_impl().await,
            controller: self,
        }
    }

    pub async fn run(
        self,
        previous_epoch_ending_ledger_info: Option<&LedgerInfo>,
    ) -> Result<Vec<LedgerInfo>> {
        self.preheat()
            .await
            .run(previous_epoch_ending_ledger_info)
            .await
    }
}

impl EpochEndingRestoreController {
    fn name(&self) -> String {
        format!("epoch ending {}", self.run_mode.name())
    }

    async fn preheat_impl(&self) -> Result<EpochEndingRestorePreheatData> {
        let manifest: EpochEndingBackup =
            self.storage.load_json_file(&self.manifest_handle).await?;
        manifest.verify()?;

        let mut next_epoch = manifest.first_epoch;
        let mut waypoint_iter = manifest.waypoints.iter();

        let mut previous_li: Option<&LedgerInfoWithSignatures> = None;
        let mut ledger_infos = Vec::new();

        let mut past_target = false;
        for chunk in &manifest.chunks {
            if past_target {
                break;
            }

            let lis = self.read_chunk(&chunk.ledger_infos).await?;
            ensure!(
                chunk.first_epoch + lis.len() as u64 == chunk.last_epoch + 1,
                "Number of items in chunks doesn't match that in manifest. \
                first_epoch: {}, last_epoch: {}, items in chunk: {}",
                chunk.first_epoch,
                chunk.last_epoch,
                lis.len(),
            );

            for li in lis {
                if li.ledger_info().version() > self.target_version {
                    past_target = true;
                    break;
                }

                ensure!(
                    li.ledger_info().epoch() == next_epoch,
                    "LedgerInfo epoch not expected. Expected: {}, actual: {}.",
                    li.ledger_info().epoch(),
                    next_epoch,
                );
                let wp_manifest = waypoint_iter.next().ok_or_else(|| {
                    anyhow!("More LedgerInfo's found than waypoints in manifest.")
                })?;
                let wp_li = Waypoint::new_epoch_boundary(li.ledger_info())?;
                ensure!(
                    *wp_manifest == wp_li,
                    "Waypoints don't match. In manifest: {}, In chunk: {}",
                    wp_manifest,
                    wp_li,
                );
                if let Some(wp_trusted) = self.trusted_waypoints.get(&wp_li.version()) {
                    ensure!(
                        *wp_trusted == wp_li,
                        "Waypoints don't match. In backup: {}, trusted: {}",
                        wp_li,
                        wp_trusted,
                    );
                } else if let Some(pre_li) = previous_li {
                    pre_li
                        .ledger_info()
                        .next_epoch_state()
                        .ok_or_else(|| {
                            anyhow!(
                                "Next epoch state not found from LI at epoch {}.",
                                pre_li.ledger_info().epoch()
                            )
                        })?
                        .verify(&li)?;
                }
                ledger_infos.push(li);
                previous_li = ledger_infos.last();
                next_epoch += 1;
            }
        }

        Ok(EpochEndingRestorePreheatData {
            manifest,
            ledger_infos,
        })
    }

    async fn read_chunk(
        &self,
        file_handle: &FileHandleRef,
    ) -> Result<Vec<LedgerInfoWithSignatures>> {
        let mut file = self.storage.open_for_read(file_handle).await?;
        let mut chunk = vec![];

        while let Some(record_bytes) = file.read_record_bytes().await? {
            chunk.push(bcs::from_bytes(&record_bytes)?);
        }

        Ok(chunk)
    }
}

struct EpochEndingRestorePreheatData {
    manifest: EpochEndingBackup,
    ledger_infos: Vec<LedgerInfoWithSignatures>,
}

pub struct PreheatedEpochEndingRestore {
    controller: EpochEndingRestoreController,
    preheat_result: Result<EpochEndingRestorePreheatData>,
}

impl PreheatedEpochEndingRestore {
    pub async fn run(
        self,
        previous_epoch_ending_ledger_info: Option<&LedgerInfo>,
    ) -> Result<Vec<LedgerInfo>> {
        let name = self.controller.name();
        info!(
            "{} started. Manifest: {}",
            name, self.controller.manifest_handle
        );
        let res = self
            .run_impl(previous_epoch_ending_ledger_info)
            .await
            .map_err(|e| anyhow!("{} failed: {}", name, e))?;
        info!("{} succeeded.", name);
        Ok(res)
    }
}

impl PreheatedEpochEndingRestore {
    async fn run_impl(
        self,
        previous_epoch_ending_ledger_info: Option<&LedgerInfo>,
    ) -> Result<Vec<LedgerInfo>> {
        let preheat_data = self
            .preheat_result
            .map_err(|e| anyhow!("Preheat failed: {}", e))?;

        let first_li = preheat_data
            .ledger_infos
            .first()
            .expect("Epoch ending backup can't be empty.");

        if let Some(li) = previous_epoch_ending_ledger_info {
            ensure!(
                li.next_block_epoch() == preheat_data.manifest.first_epoch,
                "Previous epoch ending LedgerInfo is not the one expected. \
                My first epoch: {}, previous LedgerInfo next_block_epoch: {}",
                preheat_data.manifest.first_epoch,
                li.next_block_epoch(),
            );
            // Waypoint has been verified in preheat if it's trusted, otherwise try to check
            // the signatures.
            if self
                .controller
                .trusted_waypoints
                .get(&first_li.ledger_info().version())
                .is_none()
            {
                li.next_epoch_state()
                    .ok_or_else(|| {
                        anyhow!("Previous epoch ending LedgerInfo doesn't end an epoch")
                    })?
                    .verify(first_li)?;
            }
        }

        let last_li = preheat_data
            .ledger_infos
            .last()
            .expect("Verified not empty.")
            .ledger_info();
        match self.controller.run_mode.as_ref() {
            RestoreRunMode::Restore { restore_handler } => {
                restore_handler.save_ledger_infos(&preheat_data.ledger_infos)?;

                EPOCH_ENDING_EPOCH.set(last_li.epoch() as i64);
                EPOCH_ENDING_VERSION.set(last_li.version() as i64);
            }
            RestoreRunMode::Verify => {
                VERIFY_EPOCH_ENDING_EPOCH.set(last_li.epoch() as i64);
                VERIFY_EPOCH_ENDING_VERSION.set(last_li.version() as i64);
            }
        };

        Ok(preheat_data
            .ledger_infos
            .into_iter()
            .map(|x| x.ledger_info().clone())
            .collect())
    }
}

/// Represents a history of epoch changes since epoch 0.
#[derive(Clone)]
pub struct EpochHistory {
    pub epoch_endings: Vec<LedgerInfo>,
    pub trusted_waypoints: Arc<HashMap<Version, Waypoint>>,
}

impl EpochHistory {
    pub fn verify_ledger_info(&self, li_with_sigs: &LedgerInfoWithSignatures) -> Result<()> {
        let epoch = li_with_sigs.ledger_info().epoch();
        ensure!(!self.epoch_endings.is_empty(), "Empty epoch history.",);
        if epoch > self.epoch_endings.len() as u64 {
            // TODO(aldenhu): fix this from upper level
            warn!(
                epoch = epoch,
                epoch_history_until = self.epoch_endings.len(),
                "Epoch is too new and can't be verified. Previous chunks are verified and node \
                won't be able to start if this data is malicious."
            );
            return Ok(());
        }
        if epoch == 0 {
            ensure!(
                li_with_sigs.ledger_info() == &self.epoch_endings[0],
                "Genesis epoch LedgerInfo info doesn't match.",
            );
        } else if let Some(wp_trusted) = self
            .trusted_waypoints
            .get(&li_with_sigs.ledger_info().version())
        {
            let wp_li = Waypoint::new_any(li_with_sigs.ledger_info());
            ensure!(
                *wp_trusted == wp_li,
                "Waypoints don't match. In backup: {}, trusted: {}",
                wp_li,
                wp_trusted,
            );
        } else {
            self.epoch_endings[epoch as usize - 1]
                .next_epoch_state()
                .ok_or_else(|| anyhow!("Shouldn't contain non- epoch bumping LIs."))?
                .verify(li_with_sigs)?;
        };
        Ok(())
    }
}

pub struct EpochHistoryRestoreController {
    storage: Arc<dyn BackupStorage>,
    manifest_handles: Vec<FileHandle>,
    global_opt: GlobalRestoreOptions,
}

impl EpochHistoryRestoreController {
    pub fn new(
        manifest_handles: Vec<FileHandle>,
        global_opt: GlobalRestoreOptions,
        storage: Arc<dyn BackupStorage>,
    ) -> Self {
        Self {
            storage,
            manifest_handles,
            global_opt,
        }
    }

    pub async fn run(self) -> Result<EpochHistory> {
        let name = self.name();
        info!("{} started.", name,);
        let res = self
            .run_impl()
            .await
            .map_err(|e| anyhow!("{} failed: {}", name, e))?;
        info!("{} succeeded.", name);
        Ok(res)
    }
}

impl EpochHistoryRestoreController {
    fn name(&self) -> String {
        format!("epoch history {}", self.global_opt.run_mode.name())
    }

    async fn run_impl(self) -> Result<EpochHistory> {
        let timer = Instant::now();
        if self.manifest_handles.is_empty() {
            return Ok(EpochHistory {
                epoch_endings: Vec::new(),
                trusted_waypoints: Arc::new(HashMap::new()),
            });
        }

        let futs_iter = self.manifest_handles.iter().map(|hdl| {
            EpochEndingRestoreController::new(
                EpochEndingRestoreOpt {
                    manifest_handle: hdl.clone(),
                },
                self.global_opt.clone(),
                self.storage.clone(),
            )
            .preheat()
        });
        let mut futs_stream = futures::stream::iter(futs_iter).buffered_x(
            self.global_opt.concurrent_downloads * 2, /* buffer size */
            self.global_opt.concurrent_downloads,     /* concurrency */
        );

        let mut next_epoch = 0u64;
        let mut previous_li = None;
        let mut epoch_endings = Vec::new();

        while let Some(preheated_restore) = futs_stream.next().await {
            let manifest_handle = preheated_restore.controller.manifest_handle.clone();
            let lis = preheated_restore.run(previous_li).await?;
            ensure!(
                !lis.is_empty(),
                "No epochs restored from {}",
                manifest_handle,
            );
            for li in &lis {
                ensure!(
                    li.epoch() == next_epoch,
                    "Restored LedgerInfo has epoch {}, expecting {}.",
                    li.epoch(),
                    next_epoch,
                );
                ensure!(
                    li.ends_epoch(),
                    "LedgerInfo is not one at an epoch ending. epoch: {}",
                    li.epoch(),
                );
                next_epoch += 1;
            }

            epoch_endings.extend(lis);
            previous_li = epoch_endings.last();
        }

        info!(
            "Epoch history recovered in {:.2} seconds",
            timer.elapsed().as_secs_f64()
        );
        Ok(EpochHistory {
            epoch_endings,
            trusted_waypoints: self.global_opt.trusted_waypoints.clone(),
        })
    }
}
<|endoftext|>


<commit_before>use rand;
use rand::distributions::{IndependentSample, Range};
use redis;
use redis::Value::{Nil, Okay};
use redis::{RedisResult, Value};
use std::fs::File;
use std::io::{self, Read};
use std::path::Path;
use std::thread::sleep;
use std::time::{Duration, Instant};

const DEFAULT_RETRY_COUNT: u32 = 3;
const DEFAULT_RETRY_DELAY: u32 = 200;
const CLOCK_DRIFT_FACTOR: f32 = 0.01;
const UNLOCK_SCRIPT: &str = r"if redis.call('get',KEYS[1]) == ARGV[1] then
                                return redis.call('del',KEYS[1])
                              else
                                return 0
                              end";

/// The lock manager.
///
/// Implements the necessary functionality to acquire and release locks
/// and handles the Redis connections.
pub struct RedLock {
    /// List of all Redis clients
    pub servers: Vec<redis::Client>,
    quorum: u32,
    retry_count: u32,
    retry_delay: u32,
}

pub struct Lock<'a> {
    /// The resource to lock. Will be used as the key in Redis.
    pub resource: Vec<u8>,
    /// The value for this lock.
    pub val: Vec<u8>,
    /// Time the lock is still valid.
    /// Should only be slightly smaller than the requested TTL.
    pub validity_time: usize,
    /// Used to limit the lifetime of a lock to its lock manager.
    pub lock_manager: &'a RedLock,
}

impl RedLock {
    /// Create a new lock manager instance, defined by the given Redis connection uris.
    /// Quorum is defined to be N/2+1, with N being the number of given Redis instances.
    ///
    /// Sample URI: `"redis://127.0.0.1:6379"`
    pub fn new(uris: Vec<&str>) -> RedLock {
        let quorum = (uris.len() as u32) / 2 + 1;
        let mut servers = Vec::with_capacity(uris.len());

        for &uri in uris.iter() {
            servers.push(redis::Client::open(uri).unwrap())
        }

        RedLock {
            servers,
            quorum,
            retry_count: DEFAULT_RETRY_COUNT,
            retry_delay: DEFAULT_RETRY_DELAY,
        }
    }

    /// Get 20 random bytes from `/dev/urandom`.
    pub fn get_unique_lock_id(&self) -> io::Result<Vec<u8>> {
        let file = File::open(&Path::new("/dev/urandom")).unwrap();
        let mut buf = Vec::with_capacity(20);
        match file.take(20).read_to_end(&mut buf) {
            Ok(20) => Ok(buf),
            Ok(_) => Err(io::Error::new(
                io::ErrorKind::Other,
                "Can't read enough random bytes",
            )),
            Err(e) => Err(e),
        }
    }

    /// Set retry count and retry delay.
    ///
    /// Retry count defaults to `3`.
    /// Retry delay defaults to `200`.
    pub fn set_retry(&mut self, count: u32, delay: u32) {
        self.retry_count = count;
        self.retry_delay = delay;
    }

    fn lock_instance(
        &self,
        client: &redis::Client,
        resource: &[u8],
        val: &[u8],
        ttl: usize,
    ) -> bool {
        let con = match client.get_connection() {
            Err(_) => return false,
            Ok(val) => val,
        };
        let result: RedisResult<Value> = redis::cmd("SET")
            .arg(resource)
            .arg(val)
            .arg("nx")
            .arg("px")
            .arg(ttl)
            .query(&con);
        match result {
            Ok(Okay) => true,
            Ok(Nil) => false,
            Ok(_) => false,
            Err(_) => false,
        }
    }

    /// Acquire the lock for the given resource and the requested TTL.
    ///
    /// If it succeeds, a `Lock` instance is returned,
    /// including the value and the validity time
    ///
    /// If it fails. `None` is returned.
    /// A user should retry after a short wait time.
    pub fn lock(&self, resource: &[u8], ttl: usize) -> Option<Lock> {
        let val = self.get_unique_lock_id().unwrap();

        let between = Range::new(0, self.retry_delay);
        let mut rng = rand::thread_rng();

        for _ in 0..self.retry_count {
            let mut n = 0;
            let start_time = Instant::now();
            for client in &self.servers {
                if self.lock_instance(client, resource, &val, ttl) {
                    n += 1;
                }
            }

            let drift = (ttl as f32 * CLOCK_DRIFT_FACTOR) as usize + 2;
            let elapsed = start_time.elapsed();
            let validity_time = ttl
                - drift
                - elapsed.as_secs() as usize * 1000
                - elapsed.subsec_nanos() as usize / 1_000_000;

            if n >= self.quorum && validity_time > 0 {
                return Some(Lock {
                    lock_manager: self,
                    resource: resource.to_vec(),
                    val,
                    validity_time,
                });
            } else {
                for client in &self.servers {
                    self.unlock_instance(client, resource, &val);
                }
            }

            let n = between.ind_sample(&mut rng);
            sleep(Duration::from_millis(u64::from(n)));
        }
        None
    }

    fn unlock_instance(&self, client: &redis::Client, resource: &[u8], val: &[u8]) -> bool {
        let con = match client.get_connection() {
            Err(_) => return false,
            Ok(val) => val,
        };
        let script = redis::Script::new(UNLOCK_SCRIPT);
        let result: RedisResult<i32> = script.key(resource).arg(val).invoke(&con);
        match result {
            Ok(val) => val == 1,
            Err(_) => false,
        }
    }

    /// Unlock the given lock.
    ///
    /// Unlock is best effort. It will simply try to contact all instances
    /// and remove the key.
    pub fn unlock(&self, lock: &Lock) {
        for client in &self.servers {
            self.unlock_instance(client, &lock.resource, &lock.val);
        }
    }
}

#[test]
fn test_redlock_get_unique_id() {
    let rl = RedLock::new(vec![]);

    match rl.get_unique_lock_id() {
        Ok(id) => {
            assert_eq!(20, id.len());
        }
        err => panic!("Error thrown: {:?}", err),
    }
}

#[test]
fn test_redlock_get_unique_id_uniqueness() {
    let rl = RedLock::new(vec![]);

    let id1 = rl.get_unique_lock_id().unwrap();
    let id2 = rl.get_unique_lock_id().unwrap();

    assert_eq!(20, id1.len());
    assert_eq!(20, id2.len());
    assert!(id1 != id2);
}

#[test]
fn test_redlock_valid_instance() {
    let rl = RedLock::new(vec![
        "redis://127.0.0.1:6380/",
        "redis://127.0.0.1:6381/",
        "redis://127.0.0.1:6382/",
    ]);
    assert_eq!(3, rl.servers.len());
    assert_eq!(2, rl.quorum);
}

#[test]
fn test_redlock_direct_unlock_fails() {
    let rl = RedLock::new(vec![
        "redis://127.0.0.1:6380/",
        "redis://127.0.0.1:6381/",
        "redis://127.0.0.1:6382/",
    ]);
    let key = rl.get_unique_lock_id().unwrap();

    let val = rl.get_unique_lock_id().unwrap();
    assert_eq!(false, rl.unlock_instance(&rl.servers[0], &key, &val))
}

#[test]
fn test_redlock_direct_unlock_succeeds() {
    let rl = RedLock::new(vec![
        "redis://127.0.0.1:6380/",
        "redis://127.0.0.1:6381/",
        "redis://127.0.0.1:6382/",
    ]);
    let key = rl.get_unique_lock_id().unwrap();

    let val = rl.get_unique_lock_id().unwrap();
    let con = rl.servers[0].get_connection().unwrap();
    redis::cmd("SET").arg(&*key).arg(&*val).execute(&con);

    assert_eq!(true, rl.unlock_instance(&rl.servers[0], &key, &val))
}

#[test]
fn test_redlock_direct_lock_succeeds() {
    let rl = RedLock::new(vec![
        "redis://127.0.0.1:6380/",
        "redis://127.0.0.1:6381/",
        "redis://127.0.0.1:6382/",
    ]);
    let key = rl.get_unique_lock_id().unwrap();

    let val = rl.get_unique_lock_id().unwrap();
    let con = rl.servers[0].get_connection().unwrap();

    redis::cmd("DEL").arg(&*key).execute(&con);
    assert_eq!(true, rl.lock_instance(&rl.servers[0], &*key, &*val, 1000))
}

#[test]
fn test_redlock_unlock() {
    let rl = RedLock::new(vec![
        "redis://127.0.0.1:6380/",
        "redis://127.0.0.1:6381/",
        "redis://127.0.0.1:6382/",
    ]);
    let key = rl.get_unique_lock_id().unwrap();

    let val = rl.get_unique_lock_id().unwrap();
    let con = rl.servers[0].get_connection().unwrap();
    let _: () = redis::cmd("SET").arg(&*key).arg(&*val).query(&con).unwrap();

    let lock = Lock {
        lock_manager: &rl,
        resource: key,
        val,
        validity_time: 0,
    };
    assert_eq!((), rl.unlock(&lock))
}

#[test]
fn test_redlock_lock() {
    let rl = RedLock::new(vec![
        "redis://127.0.0.1:6380/",
        "redis://127.0.0.1:6381/",
        "redis://127.0.0.1:6382/",
    ]);

    let key = rl.get_unique_lock_id().unwrap();
    match rl.lock(&key, 1000) {
        Some(lock) => {
            assert_eq!(key, lock.resource);
            assert_eq!(20, lock.val.len());
            assert!(lock.validity_time > 900);
        }
        None => panic!("Lock failed"),
    }
}

#[test]
fn test_redlock_lock_unlock() {
    let rl = RedLock::new(vec![
        "redis://127.0.0.1:6380/",
        "redis://127.0.0.1:6381/",
        "redis://127.0.0.1:6382/",
    ]);
    let rl2 = RedLock::new(vec![
        "redis://127.0.0.1:6380/",
        "redis://127.0.0.1:6381/",
        "redis://127.0.0.1:6382/",
    ]);

    let key = rl.get_unique_lock_id().unwrap();

    let lock = rl.lock(&key, 1000).unwrap();
    assert!(lock.validity_time > 900);

    match rl2.lock(&key, 1000) {
        Some(_l) => panic!("Lock acquired, even though it should be locked"),
        None => (),
    }

    rl.unlock(&lock);

    match rl2.lock(&key, 1000) {
        Some(l) => assert!(l.validity_time > 900),
        None => panic!("Lock couldn't be acquired"),
    }
}
<commit_msg>Simplify opening urandom.<commit_after>use rand;
use rand::distributions::{IndependentSample, Range};
use redis;
use redis::Value::{Nil, Okay};
use redis::{RedisResult, Value};

use std::fs::File;
use std::io::{self, Read};
use std::thread::sleep;
use std::time::{Duration, Instant};

const DEFAULT_RETRY_COUNT: u32 = 3;
const DEFAULT_RETRY_DELAY: u32 = 200;
const CLOCK_DRIFT_FACTOR: f32 = 0.01;
const UNLOCK_SCRIPT: &str = r"if redis.call('get',KEYS[1]) == ARGV[1] then
                                return redis.call('del',KEYS[1])
                              else
                                return 0
                              end";

/// The lock manager.
///
/// Implements the necessary functionality to acquire and release locks
/// and handles the Redis connections.
pub struct RedLock {
    /// List of all Redis clients
    pub servers: Vec<redis::Client>,
    quorum: u32,
    retry_count: u32,
    retry_delay: u32,
}

pub struct Lock<'a> {
    /// The resource to lock. Will be used as the key in Redis.
    pub resource: Vec<u8>,
    /// The value for this lock.
    pub val: Vec<u8>,
    /// Time the lock is still valid.
    /// Should only be slightly smaller than the requested TTL.
    pub validity_time: usize,
    /// Used to limit the lifetime of a lock to its lock manager.
    pub lock_manager: &'a RedLock,
}

impl RedLock {
    /// Create a new lock manager instance, defined by the given Redis connection uris.
    /// Quorum is defined to be N/2+1, with N being the number of given Redis instances.
    ///
    /// Sample URI: `"redis://127.0.0.1:6379"`
    pub fn new(uris: Vec<&str>) -> RedLock {
        let quorum = (uris.len() as u32) / 2 + 1;
        let mut servers = Vec::with_capacity(uris.len());

        for &uri in uris.iter() {
            servers.push(redis::Client::open(uri).unwrap())
        }

        RedLock {
            servers,
            quorum,
            retry_count: DEFAULT_RETRY_COUNT,
            retry_delay: DEFAULT_RETRY_DELAY,
        }
    }

    /// Get 20 random bytes from `/dev/urandom`.
    pub fn get_unique_lock_id(&self) -> io::Result<Vec<u8>> {
        let file = File::open("/dev/urandom")?;
        let mut buf = Vec::with_capacity(20);
        match file.take(20).read_to_end(&mut buf) {
            Ok(20) => Ok(buf),
            Ok(_) => Err(io::Error::new(
                io::ErrorKind::Other,
                "Can't read enough random bytes",
            )),
            Err(e) => Err(e),
        }
    }

    /// Set retry count and retry delay.
    ///
    /// Retry count defaults to `3`.
    /// Retry delay defaults to `200`.
    pub fn set_retry(&mut self, count: u32, delay: u32) {
        self.retry_count = count;
        self.retry_delay = delay;
    }

    fn lock_instance(
        &self,
        client: &redis::Client,
        resource: &[u8],
        val: &[u8],
        ttl: usize,
    ) -> bool {
        let con = match client.get_connection() {
            Err(_) => return false,
            Ok(val) => val,
        };
        let result: RedisResult<Value> = redis::cmd("SET")
            .arg(resource)
            .arg(val)
            .arg("nx")
            .arg("px")
            .arg(ttl)
            .query(&con);
        match result {
            Ok(Okay) => true,
            Ok(Nil) => false,
            Ok(_) => false,
            Err(_) => false,
        }
    }

    /// Acquire the lock for the given resource and the requested TTL.
    ///
    /// If it succeeds, a `Lock` instance is returned,
    /// including the value and the validity time
    ///
    /// If it fails. `None` is returned.
    /// A user should retry after a short wait time.
    pub fn lock(&self, resource: &[u8], ttl: usize) -> Option<Lock> {
        let val = self.get_unique_lock_id().unwrap();

        let between = Range::new(0, self.retry_delay);
        let mut rng = rand::thread_rng();

        for _ in 0..self.retry_count {
            let mut n = 0;
            let start_time = Instant::now();
            for client in &self.servers {
                if self.lock_instance(client, resource, &val, ttl) {
                    n += 1;
                }
            }

            let drift = (ttl as f32 * CLOCK_DRIFT_FACTOR) as usize + 2;
            let elapsed = start_time.elapsed();
            let validity_time = ttl
                - drift
                - elapsed.as_secs() as usize * 1000
                - elapsed.subsec_nanos() as usize / 1_000_000;

            if n >= self.quorum && validity_time > 0 {
                return Some(Lock {
                    lock_manager: self,
                    resource: resource.to_vec(),
                    val,
                    validity_time,
                });
            } else {
                for client in &self.servers {
                    self.unlock_instance(client, resource, &val);
                }
            }

            let n = between.ind_sample(&mut rng);
            sleep(Duration::from_millis(u64::from(n)));
        }
        None
    }

    fn unlock_instance(&self, client: &redis::Client, resource: &[u8], val: &[u8]) -> bool {
        let con = match client.get_connection() {
            Err(_) => return false,
            Ok(val) => val,
        };
        let script = redis::Script::new(UNLOCK_SCRIPT);
        let result: RedisResult<i32> = script.key(resource).arg(val).invoke(&con);
        match result {
            Ok(val) => val == 1,
            Err(_) => false,
        }
    }

    /// Unlock the given lock.
    ///
    /// Unlock is best effort. It will simply try to contact all instances
    /// and remove the key.
    pub fn unlock(&self, lock: &Lock) {
        for client in &self.servers {
            self.unlock_instance(client, &lock.resource, &lock.val);
        }
    }
}

#[test]
fn test_redlock_get_unique_id() {
    let rl = RedLock::new(vec![]);

    match rl.get_unique_lock_id() {
        Ok(id) => {
            assert_eq!(20, id.len());
        }
        err => panic!("Error thrown: {:?}", err),
    }
}

#[test]
fn test_redlock_get_unique_id_uniqueness() {
    let rl = RedLock::new(vec![]);

    let id1 = rl.get_unique_lock_id().unwrap();
    let id2 = rl.get_unique_lock_id().unwrap();

    assert_eq!(20, id1.len());
    assert_eq!(20, id2.len());
    assert!(id1 != id2);
}

#[test]
fn test_redlock_valid_instance() {
    let rl = RedLock::new(vec![
        "redis://127.0.0.1:6380/",
        "redis://127.0.0.1:6381/",
        "redis://127.0.0.1:6382/",
    ]);
    assert_eq!(3, rl.servers.len());
    assert_eq!(2, rl.quorum);
}

#[test]
fn test_redlock_direct_unlock_fails() {
    let rl = RedLock::new(vec![
        "redis://127.0.0.1:6380/",
        "redis://127.0.0.1:6381/",
        "redis://127.0.0.1:6382/",
    ]);
    let key = rl.get_unique_lock_id().unwrap();

    let val = rl.get_unique_lock_id().unwrap();
    assert_eq!(false, rl.unlock_instance(&rl.servers[0], &key, &val))
}

#[test]
fn test_redlock_direct_unlock_succeeds() {
    let rl = RedLock::new(vec![
        "redis://127.0.0.1:6380/",
        "redis://127.0.0.1:6381/",
        "redis://127.0.0.1:6382/",
    ]);
    let key = rl.get_unique_lock_id().unwrap();

    let val = rl.get_unique_lock_id().unwrap();
    let con = rl.servers[0].get_connection().unwrap();
    redis::cmd("SET").arg(&*key).arg(&*val).execute(&con);

    assert_eq!(true, rl.unlock_instance(&rl.servers[0], &key, &val))
}

#[test]
fn test_redlock_direct_lock_succeeds() {
    let rl = RedLock::new(vec![
        "redis://127.0.0.1:6380/",
        "redis://127.0.0.1:6381/",
        "redis://127.0.0.1:6382/",
    ]);
    let key = rl.get_unique_lock_id().unwrap();

    let val = rl.get_unique_lock_id().unwrap();
    let con = rl.servers[0].get_connection().unwrap();

    redis::cmd("DEL").arg(&*key).execute(&con);
    assert_eq!(true, rl.lock_instance(&rl.servers[0], &*key, &*val, 1000))
}

#[test]
fn test_redlock_unlock() {
    let rl = RedLock::new(vec![
        "redis://127.0.0.1:6380/",
        "redis://127.0.0.1:6381/",
        "redis://127.0.0.1:6382/",
    ]);
    let key = rl.get_unique_lock_id().unwrap();

    let val = rl.get_unique_lock_id().unwrap();
    let con = rl.servers[0].get_connection().unwrap();
    let _: () = redis::cmd("SET").arg(&*key).arg(&*val).query(&con).unwrap();

    let lock = Lock {
        lock_manager: &rl,
        resource: key,
        val,
        validity_time: 0,
    };
    assert_eq!((), rl.unlock(&lock))
}

#[test]
fn test_redlock_lock() {
    let rl = RedLock::new(vec![
        "redis://127.0.0.1:6380/",
        "redis://127.0.0.1:6381/",
        "redis://127.0.0.1:6382/",
    ]);

    let key = rl.get_unique_lock_id().unwrap();
    match rl.lock(&key, 1000) {
        Some(lock) => {
            assert_eq!(key, lock.resource);
            assert_eq!(20, lock.val.len());
            assert!(lock.validity_time > 900);
        }
        None => panic!("Lock failed"),
    }
}

#[test]
fn test_redlock_lock_unlock() {
    let rl = RedLock::new(vec![
        "redis://127.0.0.1:6380/",
        "redis://127.0.0.1:6381/",
        "redis://127.0.0.1:6382/",
    ]);
    let rl2 = RedLock::new(vec![
        "redis://127.0.0.1:6380/",
        "redis://127.0.0.1:6381/",
        "redis://127.0.0.1:6382/",
    ]);

    let key = rl.get_unique_lock_id().unwrap();

    let lock = rl.lock(&key, 1000).unwrap();
    assert!(lock.validity_time > 900);

    match rl2.lock(&key, 1000) {
        Some(_l) => panic!("Lock acquired, even though it should be locked"),
        None => (),
    }

    rl.unlock(&lock);

    match rl2.lock(&key, 1000) {
        Some(l) => assert!(l.validity_time > 900),
        None => panic!("Lock couldn't be acquired"),
    }
}
<|endoftext|>


<commit_before>use std::fmt::Write;

use crate::{Assist, AssistId, AssistCtx};

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

use itertools::Itertools;

enum AddMissingImplMembersMode {
    DefaultMethodsOnly,
    NoDefaultMethods,
}

pub(crate) fn add_missing_impl_members(ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    add_missing_impl_members_inner(
        ctx,
        AddMissingImplMembersMode::NoDefaultMethods,
        "add_impl_missing_members",
        "add missing impl members",
    )
}

pub(crate) fn add_missing_default_members(ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    add_missing_impl_members_inner(
        ctx,
        AddMissingImplMembersMode::DefaultMethodsOnly,
        "add_impl_default_members",
        "add impl default members",
    )
}

fn add_missing_impl_members_inner(
    mut ctx: AssistCtx<impl HirDatabase>,
    mode: AddMissingImplMembersMode,
    assist_id: &'static str,
    label: &'static str,
) -> 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 analyzer = hir::SourceAnalyzer::new(ctx.db, position.file_id, impl_node.syntax(), None);

        resolve_target_trait_def(ctx.db, &analyzer, 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| match mode {
                AddMissingImplMembersMode::DefaultMethodsOnly => t.body().is_some(),
                AddMissingImplMembersMode::NoDefaultMethods => t.body().is_none(),
            })
            .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(assist_id), label, |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_with_tokens();
            let last_whitespace =
                children.filter_map(|it| ast::Whitespace::cast(it.as_token()?)).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,
    analyzer: &hir::SourceAnalyzer,
    impl_block: &ast::ImplBlock,
) -> Option<TreeArc<ast::TraitDef>> {
    let ast_path =
        impl_block.target_trait().map(AstNode::syntax).and_then(ast::PathType::cast)?.path()?;

    match analyzer.resolve_path(db, &ast_path) {
        Some(hir::PathResolution::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_with_tokens() {
        match (child.prev_sibling_or_token().map(|c| c.kind()), child.kind()) {
            (_, SyntaxKind::SEMI) => buf.push_str(" { unimplemented!() }"),
            (_, SyntaxKind::ATTR) | (_, SyntaxKind::COMMENT) => {}
            (Some(SyntaxKind::ATTR), SyntaxKind::WHITESPACE)
            | (Some(SyntaxKind::COMMENT), SyntaxKind::WHITESPACE) => {}
            _ => write!(buf, "{}", child).unwrap(),
        };
    }

    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!() }<|>
}",
        );
    }

    #[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_and_default_methods() {
        check_assist_not_applicable(
            add_missing_impl_members,
            "
trait Foo {
    fn (arg: u32);
    fn valid(some: u32) -> bool { false }
}
struct S;
impl Foo for S { <|> }",
        )
    }

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

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

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

    #[test]
    fn test_with_docstring_and_attrs() {
        check_assist(
            add_missing_impl_members,
            r#"
#[doc(alias = "test alias")]
trait Foo {
    /// doc string
    #[must_use]
    fn foo(&self);
}
struct S;
impl Foo for S {}<|>"#,
            r#"
#[doc(alias = "test alias")]
trait Foo {
    /// doc string
    #[must_use]
    fn foo(&self);
}
struct S;
impl Foo for S {
    fn foo(&self) { unimplemented!() }<|>
}"#,
        )
    }

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

}
<commit_msg>better formatting when adding trait members<commit_after>use std::fmt::Write;

use crate::{Assist, AssistId, AssistCtx};

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

use itertools::Itertools;

enum AddMissingImplMembersMode {
    DefaultMethodsOnly,
    NoDefaultMethods,
}

pub(crate) fn add_missing_impl_members(ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    add_missing_impl_members_inner(
        ctx,
        AddMissingImplMembersMode::NoDefaultMethods,
        "add_impl_missing_members",
        "add missing impl members",
    )
}

pub(crate) fn add_missing_default_members(ctx: AssistCtx<impl HirDatabase>) -> Option<Assist> {
    add_missing_impl_members_inner(
        ctx,
        AddMissingImplMembersMode::DefaultMethodsOnly,
        "add_impl_default_members",
        "add impl default members",
    )
}

fn add_missing_impl_members_inner(
    mut ctx: AssistCtx<impl HirDatabase>,
    mode: AddMissingImplMembersMode,
    assist_id: &'static str,
    label: &'static str,
) -> 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 analyzer = hir::SourceAnalyzer::new(ctx.db, position.file_id, impl_node.syntax(), None);

        resolve_target_trait_def(ctx.db, &analyzer, 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| match mode {
                AddMissingImplMembersMode::DefaultMethodsOnly => t.body().is_some(),
                AddMissingImplMembersMode::NoDefaultMethods => t.body().is_none(),
            })
            .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(assist_id), label, |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_with_tokens();
            let last_whitespace =
                children.filter_map(|it| ast::Whitespace::cast(it.as_token()?)).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);
        // FIXME: place the cursor on the first unimplemented?
        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,
    analyzer: &hir::SourceAnalyzer,
    impl_block: &ast::ImplBlock,
) -> Option<TreeArc<ast::TraitDef>> {
    let ast_path =
        impl_block.target_trait().map(AstNode::syntax).and_then(ast::PathType::cast)?.path()?;

    match analyzer.resolve_path(db, &ast_path) {
        Some(hir::PathResolution::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_with_tokens() {
        match (child.prev_sibling_or_token().map(|c| c.kind()), child.kind()) {
            (_, SyntaxKind::SEMI) => buf.push_str(" {\n    unimplemented!()\n}"),
            (_, SyntaxKind::ATTR) | (_, SyntaxKind::COMMENT) => {}
            (Some(SyntaxKind::ATTR), SyntaxKind::WHITESPACE)
            | (Some(SyntaxKind::COMMENT), SyntaxKind::WHITESPACE) => {}
            _ => write!(buf, "{}", child).unwrap(),
        };
    }

    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!()
    }<|>
}",
        );
    }

    #[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_and_default_methods() {
        check_assist_not_applicable(
            add_missing_impl_members,
            "
trait Foo {
    fn (arg: u32);
    fn valid(some: u32) -> bool { false }
}
struct S;
impl Foo for S { <|> }",
        )
    }

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

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

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

    #[test]
    fn test_with_docstring_and_attrs() {
        check_assist(
            add_missing_impl_members,
            r#"
#[doc(alias = "test alias")]
trait Foo {
    /// doc string
    #[must_use]
    fn foo(&self);
}
struct S;
impl Foo for S {}<|>"#,
            r#"
#[doc(alias = "test alias")]
trait Foo {
    /// doc string
    #[must_use]
    fn foo(&self);
}
struct S;
impl Foo for S {
    fn foo(&self) {
        unimplemented!()
    }<|>
}"#,
        )
    }

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

}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! A windowing implementation using GLFW.

use windowing::{ApplicationMethods, WindowEvent, WindowMethods};
use windowing::{IdleWindowEvent, ResizeWindowEvent, LoadUrlWindowEvent, MouseWindowEventClass};
use windowing::{ScrollWindowEvent, ZoomWindowEvent, NavigationWindowEvent, FinishedWindowEvent};
use windowing::{QuitWindowEvent, MouseWindowClickEvent, MouseWindowMouseDownEvent, MouseWindowMouseUpEvent};
use windowing::{Forward, Back};

use alert::{Alert, AlertMethods};
use std::libc::c_int;
use std::local_data;
use geom::point::Point2D;
use geom::size::Size2D;
use servo_msg::compositor_msg::{IdleRenderState, RenderState, RenderingRenderState};
use servo_msg::compositor_msg::{FinishedLoading, Blank, Loading, PerformingLayout, ReadyState};

use glfw;

static THROBBER: [char, ..8] = [ '⣾', '⣽', '⣻', '⢿', '⡿', '⣟', '⣯', '⣷' ];

/// A structure responsible for setting up and tearing down the entire windowing system.
pub struct Application;

impl ApplicationMethods for Application {
    fn new() -> Application {
        glfw::init();
        Application
    }
}

impl Drop for Application {
    fn drop(&self) {
        glfw::terminate();
        drop_local_window();
    }
}

/// The type of a window.
pub struct Window {
    glfw_window: glfw::Window,

    event_queue: @mut ~[WindowEvent],

    drag_origin: Point2D<c_int>,

    mouse_down_button: @mut c_int,
    mouse_down_point: @mut Point2D<c_int>,

    ready_state: ReadyState,
    render_state: RenderState,
    throbber_frame: u8,
}

impl WindowMethods<Application> for Window {
    /// Creates a new window.
    fn new(_: &Application) -> @mut Window {
        // Create the GLFW window.
        let glfw_window = glfw::Window::create(800, 600, "Servo", glfw::Windowed)
            .expect("Failed to create GLFW window");
        glfw_window.make_context_current();

        // Create our window object.
        let window = @mut Window {
            glfw_window: glfw_window,

            event_queue: @mut ~[],

            drag_origin: Point2D(0 as c_int, 0),

            mouse_down_button: @mut 0,
            mouse_down_point: @mut Point2D(0 as c_int, 0),

            ready_state: Blank,
            render_state: IdleRenderState,
            throbber_frame: 0,
        };

        install_local_window(window);

        // Register event handlers.
        do window.glfw_window.set_framebuffer_size_callback |_win, width, height| {
            local_window().event_queue.push(ResizeWindowEvent(width as uint, height as uint))
        }
        do window.glfw_window.set_key_callback |_win, key, _scancode, action, mods| {
            if action == glfw::PRESS {
                local_window().handle_key(key, mods)
            }
        }
        do window.glfw_window.set_mouse_button_callback |win, button, action, _mods| {
            let (x, y) = win.get_cursor_pos();
            //handle hidpi displays, since GLFW returns non-hi-def coordinates.
            let (backing_size, _) = win.get_framebuffer_size();
            let (window_size, _) = win.get_size();
            let hidpi = (backing_size as f32) / (window_size as f32);
            let x = x as f32 * hidpi;
            let y = y as f32 * hidpi;
            if button < 3 {
                local_window().handle_mouse(button, action, x as i32, y as i32);
            }
        }
        do window.glfw_window.set_scroll_callback |win, x_offset, y_offset| {
            let dx = (x_offset as f32) * 30.0;
            let dy = (y_offset as f32) * 30.0;

            let (x, y) = win.get_cursor_pos();
            //handle hidpi displays, since GLFW returns non-hi-def coordinates.
            let (backing_size, _) = win.get_framebuffer_size();
            let (window_size, _) = win.get_size();
            let hidpi = (backing_size as f32) / (window_size as f32);
            let x = x as f32 * hidpi;
            let y = y as f32 * hidpi;

            local_window().event_queue.push(ScrollWindowEvent(Point2D(dx, dy), Point2D(x as i32, y as i32)));
        }

        window
    }

    /// Returns the size of the window.
    fn size(&self) -> Size2D<f32> {
        let (width, height) = self.glfw_window.get_framebuffer_size();
        Size2D(width as f32, height as f32)
    }

    /// Presents the window to the screen (perhaps by page flipping).
    fn present(&mut self) {
        self.glfw_window.swap_buffers();
    }

    fn recv(@mut self) -> WindowEvent {
        if !self.event_queue.is_empty() {
            return self.event_queue.shift()
        }
        glfw::poll_events();
        self.throbber_frame = (self.throbber_frame + 1) % (THROBBER.len() as u8);
        self.update_window_title(false);
        if self.glfw_window.should_close() {
            QuitWindowEvent
        } else if !self.event_queue.is_empty() {
            self.event_queue.shift()
        } else {
            IdleWindowEvent
        }
    }

    /// Sets the ready state.
    fn set_ready_state(@mut self, ready_state: ReadyState) {
        self.ready_state = ready_state;
        self.update_window_title(true)
    }

    /// Sets the render state.
    fn set_render_state(@mut self, render_state: RenderState) {
        if self.ready_state == FinishedLoading &&
            self.render_state == RenderingRenderState &&
            render_state == IdleRenderState {
            // page loaded
            self.event_queue.push(FinishedWindowEvent);
        }

        self.render_state = render_state;
        self.update_window_title(true)
    }

    fn hidpi_factor(@mut self) -> f32 {
        let (backing_size, _) = self.glfw_window.get_framebuffer_size();
        let (window_size, _) = self.glfw_window.get_size();
        (backing_size as f32) / (window_size as f32)
    }
}

impl Window {
    /// Helper function to set the window title in accordance with the ready state.
    fn update_window_title(&self, state_change: bool) {
        let throbber = THROBBER[self.throbber_frame];
        match self.ready_state {
            Blank => {
                if state_change { self.glfw_window.set_title(fmt!("blank — Servo")) }
            }
            Loading => {
                self.glfw_window.set_title(fmt!("%c Loading — Servo", throbber))
            }
            PerformingLayout => {
                self.glfw_window.set_title(fmt!("%c Performing Layout — Servo", throbber))
            }
            FinishedLoading => {
                match self.render_state {
                    RenderingRenderState => {
                        self.glfw_window.set_title(fmt!("%c Rendering — Servo", throbber))
                    }
                    IdleRenderState => {
                        if state_change { self.glfw_window.set_title("Servo") }
                    }
                }
            }
        }
    }

    /// Helper function to handle keyboard events.
    fn handle_key(&self, key: c_int, mods: glfw::KeyMods) {
        let mods = *mods;
        match key {
            glfw::KEY_ESCAPE => self.glfw_window.set_should_close(true),
            glfw::KEY_L if mods & glfw::MOD_CONTROL != 0 => self.load_url(), // Ctrl+L
            glfw::KEY_EQUAL if mods & glfw::MOD_CONTROL != 0 => { // Ctrl-+
                self.event_queue.push(ZoomWindowEvent(1.1));
            }
            glfw::KEY_MINUS if mods & glfw::MOD_CONTROL != 0 => { // Ctrl--
                self.event_queue.push(ZoomWindowEvent(0.90909090909));
            }
            glfw::KEY_BACKSPACE if mods & glfw::MOD_SHIFT != 0 => { // Shift-Backspace
                self.event_queue.push(NavigationWindowEvent(Forward));
            }
            glfw::KEY_BACKSPACE => { // Backspace
                self.event_queue.push(NavigationWindowEvent(Back));
            }
            _ => {}
        }
    }

    /// Helper function to handle a click
    fn handle_mouse(&self, button: c_int, action: c_int, x: c_int, y: c_int) {
        // FIXME(tkuehn): max pixel dist should be based on pixel density
        let max_pixel_dist = 10f;
        let event = match action {
            glfw::PRESS => {
                *self.mouse_down_point = Point2D(x, y);
                *self.mouse_down_button = button;
                MouseWindowMouseDownEvent(button as uint, Point2D(x as f32, y as f32))
            }
            glfw::RELEASE => {
                if *self.mouse_down_button == button {
                    let pixel_dist = *self.mouse_down_point - Point2D(x, y);
                    let pixel_dist = ((pixel_dist.x * pixel_dist.x +
                                       pixel_dist.y * pixel_dist.y) as float).sqrt();
                    if pixel_dist < max_pixel_dist {
                        let click_event = MouseWindowClickEvent(button as uint,
                                                           Point2D(x as f32, y as f32));
                        self.event_queue.push(MouseWindowEventClass(click_event));
                    }
                }
                MouseWindowMouseUpEvent(button as uint, Point2D(x as f32, y as f32))
            }
            _ => fail!("I cannot recognize the type of mouse action that occured. :-(")
        };
        self.event_queue.push(MouseWindowEventClass(event));
    }

    /// Helper function to pop up an alert box prompting the user to load a URL.
    fn load_url(&self) {
        let mut alert: Alert = AlertMethods::new("Navigate to:");
        alert.add_prompt();
        alert.run();
        let value = alert.prompt_value();
        if "" == value {    // To avoid crashing on Linux.
            self.event_queue.push(LoadUrlWindowEvent(~"http://purple.com/"))
        } else {
            self.event_queue.push(LoadUrlWindowEvent(value))
        }
    }
}

static TLS_KEY: local_data::Key<@mut Window> = &local_data::Key;

fn install_local_window(window: @mut Window) {
    local_data::set(TLS_KEY, window);
}

fn drop_local_window() {
    local_data::pop(TLS_KEY);
}

fn local_window() -> @mut Window {
    local_data::get(TLS_KEY, |v| *v.unwrap())
}
<commit_msg>Print GLFW errors<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! A windowing implementation using GLFW.

use windowing::{ApplicationMethods, WindowEvent, WindowMethods};
use windowing::{IdleWindowEvent, ResizeWindowEvent, LoadUrlWindowEvent, MouseWindowEventClass};
use windowing::{ScrollWindowEvent, ZoomWindowEvent, NavigationWindowEvent, FinishedWindowEvent};
use windowing::{QuitWindowEvent, MouseWindowClickEvent, MouseWindowMouseDownEvent, MouseWindowMouseUpEvent};
use windowing::{Forward, Back};

use alert::{Alert, AlertMethods};
use std::libc::c_int;
use std::local_data;
use geom::point::Point2D;
use geom::size::Size2D;
use servo_msg::compositor_msg::{IdleRenderState, RenderState, RenderingRenderState};
use servo_msg::compositor_msg::{FinishedLoading, Blank, Loading, PerformingLayout, ReadyState};

use glfw;

static THROBBER: [char, ..8] = [ '⣾', '⣽', '⣻', '⢿', '⡿', '⣟', '⣯', '⣷' ];

/// A structure responsible for setting up and tearing down the entire windowing system.
pub struct Application;

impl ApplicationMethods for Application {
    fn new() -> Application {
        // Per GLFW docs it's safe to set the error callback before calling
        // glfwInit(), and this way we notice errors from init too.
        do glfw::set_error_callback |_error_code, description| {
            error!("GLFW error: %s", description);
        };
        glfw::init();
        Application
    }
}

impl Drop for Application {
    fn drop(&self) {
        glfw::terminate();
        drop_local_window();
    }
}

/// The type of a window.
pub struct Window {
    glfw_window: glfw::Window,

    event_queue: @mut ~[WindowEvent],

    drag_origin: Point2D<c_int>,

    mouse_down_button: @mut c_int,
    mouse_down_point: @mut Point2D<c_int>,

    ready_state: ReadyState,
    render_state: RenderState,
    throbber_frame: u8,
}

impl WindowMethods<Application> for Window {
    /// Creates a new window.
    fn new(_: &Application) -> @mut Window {
        // Create the GLFW window.
        let glfw_window = glfw::Window::create(800, 600, "Servo", glfw::Windowed)
            .expect("Failed to create GLFW window");
        glfw_window.make_context_current();

        // Create our window object.
        let window = @mut Window {
            glfw_window: glfw_window,

            event_queue: @mut ~[],

            drag_origin: Point2D(0 as c_int, 0),

            mouse_down_button: @mut 0,
            mouse_down_point: @mut Point2D(0 as c_int, 0),

            ready_state: Blank,
            render_state: IdleRenderState,
            throbber_frame: 0,
        };

        install_local_window(window);

        // Register event handlers.
        do window.glfw_window.set_framebuffer_size_callback |_win, width, height| {
            local_window().event_queue.push(ResizeWindowEvent(width as uint, height as uint))
        }
        do window.glfw_window.set_key_callback |_win, key, _scancode, action, mods| {
            if action == glfw::PRESS {
                local_window().handle_key(key, mods)
            }
        }
        do window.glfw_window.set_mouse_button_callback |win, button, action, _mods| {
            let (x, y) = win.get_cursor_pos();
            //handle hidpi displays, since GLFW returns non-hi-def coordinates.
            let (backing_size, _) = win.get_framebuffer_size();
            let (window_size, _) = win.get_size();
            let hidpi = (backing_size as f32) / (window_size as f32);
            let x = x as f32 * hidpi;
            let y = y as f32 * hidpi;
            if button < 3 {
                local_window().handle_mouse(button, action, x as i32, y as i32);
            }
        }
        do window.glfw_window.set_scroll_callback |win, x_offset, y_offset| {
            let dx = (x_offset as f32) * 30.0;
            let dy = (y_offset as f32) * 30.0;

            let (x, y) = win.get_cursor_pos();
            //handle hidpi displays, since GLFW returns non-hi-def coordinates.
            let (backing_size, _) = win.get_framebuffer_size();
            let (window_size, _) = win.get_size();
            let hidpi = (backing_size as f32) / (window_size as f32);
            let x = x as f32 * hidpi;
            let y = y as f32 * hidpi;

            local_window().event_queue.push(ScrollWindowEvent(Point2D(dx, dy), Point2D(x as i32, y as i32)));
        }

        window
    }

    /// Returns the size of the window.
    fn size(&self) -> Size2D<f32> {
        let (width, height) = self.glfw_window.get_framebuffer_size();
        Size2D(width as f32, height as f32)
    }

    /// Presents the window to the screen (perhaps by page flipping).
    fn present(&mut self) {
        self.glfw_window.swap_buffers();
    }

    fn recv(@mut self) -> WindowEvent {
        if !self.event_queue.is_empty() {
            return self.event_queue.shift()
        }
        glfw::poll_events();
        self.throbber_frame = (self.throbber_frame + 1) % (THROBBER.len() as u8);
        self.update_window_title(false);
        if self.glfw_window.should_close() {
            QuitWindowEvent
        } else if !self.event_queue.is_empty() {
            self.event_queue.shift()
        } else {
            IdleWindowEvent
        }
    }

    /// Sets the ready state.
    fn set_ready_state(@mut self, ready_state: ReadyState) {
        self.ready_state = ready_state;
        self.update_window_title(true)
    }

    /// Sets the render state.
    fn set_render_state(@mut self, render_state: RenderState) {
        if self.ready_state == FinishedLoading &&
            self.render_state == RenderingRenderState &&
            render_state == IdleRenderState {
            // page loaded
            self.event_queue.push(FinishedWindowEvent);
        }

        self.render_state = render_state;
        self.update_window_title(true)
    }

    fn hidpi_factor(@mut self) -> f32 {
        let (backing_size, _) = self.glfw_window.get_framebuffer_size();
        let (window_size, _) = self.glfw_window.get_size();
        (backing_size as f32) / (window_size as f32)
    }
}

impl Window {
    /// Helper function to set the window title in accordance with the ready state.
    fn update_window_title(&self, state_change: bool) {
        let throbber = THROBBER[self.throbber_frame];
        match self.ready_state {
            Blank => {
                if state_change { self.glfw_window.set_title(fmt!("blank — Servo")) }
            }
            Loading => {
                self.glfw_window.set_title(fmt!("%c Loading — Servo", throbber))
            }
            PerformingLayout => {
                self.glfw_window.set_title(fmt!("%c Performing Layout — Servo", throbber))
            }
            FinishedLoading => {
                match self.render_state {
                    RenderingRenderState => {
                        self.glfw_window.set_title(fmt!("%c Rendering — Servo", throbber))
                    }
                    IdleRenderState => {
                        if state_change { self.glfw_window.set_title("Servo") }
                    }
                }
            }
        }
    }

    /// Helper function to handle keyboard events.
    fn handle_key(&self, key: c_int, mods: glfw::KeyMods) {
        let mods = *mods;
        match key {
            glfw::KEY_ESCAPE => self.glfw_window.set_should_close(true),
            glfw::KEY_L if mods & glfw::MOD_CONTROL != 0 => self.load_url(), // Ctrl+L
            glfw::KEY_EQUAL if mods & glfw::MOD_CONTROL != 0 => { // Ctrl-+
                self.event_queue.push(ZoomWindowEvent(1.1));
            }
            glfw::KEY_MINUS if mods & glfw::MOD_CONTROL != 0 => { // Ctrl--
                self.event_queue.push(ZoomWindowEvent(0.90909090909));
            }
            glfw::KEY_BACKSPACE if mods & glfw::MOD_SHIFT != 0 => { // Shift-Backspace
                self.event_queue.push(NavigationWindowEvent(Forward));
            }
            glfw::KEY_BACKSPACE => { // Backspace
                self.event_queue.push(NavigationWindowEvent(Back));
            }
            _ => {}
        }
    }

    /// Helper function to handle a click
    fn handle_mouse(&self, button: c_int, action: c_int, x: c_int, y: c_int) {
        // FIXME(tkuehn): max pixel dist should be based on pixel density
        let max_pixel_dist = 10f;
        let event = match action {
            glfw::PRESS => {
                *self.mouse_down_point = Point2D(x, y);
                *self.mouse_down_button = button;
                MouseWindowMouseDownEvent(button as uint, Point2D(x as f32, y as f32))
            }
            glfw::RELEASE => {
                if *self.mouse_down_button == button {
                    let pixel_dist = *self.mouse_down_point - Point2D(x, y);
                    let pixel_dist = ((pixel_dist.x * pixel_dist.x +
                                       pixel_dist.y * pixel_dist.y) as float).sqrt();
                    if pixel_dist < max_pixel_dist {
                        let click_event = MouseWindowClickEvent(button as uint,
                                                           Point2D(x as f32, y as f32));
                        self.event_queue.push(MouseWindowEventClass(click_event));
                    }
                }
                MouseWindowMouseUpEvent(button as uint, Point2D(x as f32, y as f32))
            }
            _ => fail!("I cannot recognize the type of mouse action that occured. :-(")
        };
        self.event_queue.push(MouseWindowEventClass(event));
    }

    /// Helper function to pop up an alert box prompting the user to load a URL.
    fn load_url(&self) {
        let mut alert: Alert = AlertMethods::new("Navigate to:");
        alert.add_prompt();
        alert.run();
        let value = alert.prompt_value();
        if "" == value {    // To avoid crashing on Linux.
            self.event_queue.push(LoadUrlWindowEvent(~"http://purple.com/"))
        } else {
            self.event_queue.push(LoadUrlWindowEvent(value))
        }
    }
}

static TLS_KEY: local_data::Key<@mut Window> = &local_data::Key;

fn install_local_window(window: @mut Window) {
    local_data::set(TLS_KEY, window);
}

fn drop_local_window() {
    local_data::pop(TLS_KEY);
}

fn local_window() -> @mut Window {
    local_data::get(TLS_KEY, |v| *v.unwrap())
}
<|endoftext|>