<commit_before>// Copyright lowRISC contributors.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

use anyhow::{bail, Result};
use regex::Regex;
use std::matches;
use std::ops::Drop;
use std::thread;
use std::time::Duration;
use structopt::StructOpt;

use opentitanlib::app::TransportWrapper;
use opentitanlib::execute_test;
use opentitanlib::spiflash::{
    sfdp, BlockEraseSize, SpiFlash, SupportedAddressModes, WriteGranularity,
};
use opentitanlib::test_utils::init::InitializeTest;
use opentitanlib::uart::console::{ExitStatus, UartConsole};
use opentitanlib::util::parse_int::ParseInt;

#[derive(Debug, StructOpt)]
struct Opts {
    #[structopt(flatten)]
    init: InitializeTest,

    #[structopt(
        long, parse(try_from_str=humantime::parse_duration),
        default_value = "5s",
        help = "Bootstrap detection timeout",
    )]
    timeout: Duration,

    #[structopt(
        long, parse(try_from_str=usize::from_str),
        default_value = "12",
        help = "JEDEC page of manufacturer",
    )]
    jedec_page: usize,

    #[structopt(
        long, parse(try_from_str=u8::from_str),
        default_value = "0xEF",
        help = "JEDEC ID of manufacturer",
    )]
    jedec_id: u8,

    #[structopt(
        long, parse(try_from_str=u8::from_str),
        default_value = "0x08",
        help = "JEDEC manufacturer product ID",
    )]
    jedec_product: u8,

    #[structopt(
        long, parse(try_from_str=u32::from_str),
        default_value = "0x100000",
        help = "Size of the internal flash",
    )]
    flash_size: u32,
}

/// A struct for isolating individual test points.
///
/// This is basically an example of the RAII technique.
/// Calling `BootstrapTest::start()` resets the chip and puts it in bootstrap mode (assuming it's
/// enabled). Strapping pins are released and the chip is reset when the returned struct is
/// dropped, typically at the end of the test point.
struct BootstrapTest<'a> {
    transport: &'a TransportWrapper,
    reset_delay: Duration,
}

impl<'a> BootstrapTest<'a> {
    pub fn start(transport: &'a TransportWrapper, reset_delay: Duration) -> Result<Self> {
        let b = BootstrapTest {
            transport,
            reset_delay,
        };
        transport.apply_pin_strapping("ROM_BOOTSTRAP")?;
        transport.reset_target(b.reset_delay, true)?;

        let spi = transport.spi("0")?;
        while SpiFlash::read_status(&*spi)? != 0 {
            thread::sleep(Duration::from_millis(10));
        }

        Ok(b)
    }
}

impl<'a> Drop for BootstrapTest<'a> {
    fn drop(&mut self) {
        self.transport
            .remove_pin_strapping("ROM_BOOTSTRAP")
            .unwrap();
        self.transport.reset_target(self.reset_delay, true).unwrap();
    }
}

fn test_bootstrap_enabled_requested(opts: &Opts, transport: &TransportWrapper) -> Result<()> {
    // TODO: Should really `opts.init.uart_params.create()` here, but we need to refactor
    // BootstrapOptions first.
    //let uart = opts.init.uart_params.create(&transport)?;
    let uart = transport.uart("0")?;
    let mut console = UartConsole {
        timeout: Some(opts.timeout),
        exit_success: Some(Regex::new(r"bootstrap:1\r\n")?),
        exit_failure: Some(Regex::new(r"BFV:")?),
        ..Default::default()
    };

    transport.apply_pin_strapping("ROM_BOOTSTRAP")?;
    transport.reset_target(opts.init.bootstrap.options.reset_delay, true)?;

    // Now watch the console for the exit conditions.
    let result = console.interact(&*uart, None, Some(&mut std::io::stdout()))?;
    match result {
        ExitStatus::ExitSuccess => {}
        _ => {
            bail!("FAIL: {:?}", result);
        }
    };
    // Now check whether the SPI device is responding to status messages
    let spi = transport.spi("0")?;
    assert_eq!(SpiFlash::read_status(&*spi)?, 0x00);

    Ok(())
}

fn test_bootstrap_enabled_not_requested(opts: &Opts, transport: &TransportWrapper) -> Result<()> {
    // TODO: Should really `opts.init.uart_params.create()` here, but we need to refactor
    // BootstrapOptions first.
    //let uart = opts.init.uart_params.create(&transport)?;
    let uart = transport.uart("0")?;
    let mut console = UartConsole {
        timeout: Some(opts.timeout),
        exit_success: Some(Regex::new(r"BFV:")?),
        exit_failure: Some(Regex::new(r"bootstrap:1\r\n")?),
        ..Default::default()
    };

    transport.remove_pin_strapping("ROM_BOOTSTRAP")?;
    transport.reset_target(opts.init.bootstrap.options.reset_delay, true)?;

    // Now watch the console for the exit conditions.
    let result = console.interact(&*uart, None, Some(&mut std::io::stdout()))?;
    match result {
        ExitStatus::ExitSuccess => {}
        _ => {
            bail!("FAIL: {:?}", result);
        }
    };

    // Now check whether the SPI device is responding to commands.
    // Note: CIPO line is in high-z state when CMD_INFO registers are not configured.
    // Use READ_SFDP instead of READ_STATUS to avoid false negatives when bootstrap is not
    // requested
    let spi = transport.spi("0")?;
    assert!(matches!(
        SpiFlash::read_sfdp(&*spi)
            .unwrap_err()
            .downcast::<sfdp::Error>()
            .unwrap(),
        sfdp::Error::WrongHeaderSignature(..)
    ));
    Ok(())
}

fn test_jedec_id(opts: &Opts, transport: &TransportWrapper) -> Result<()> {
    let spi = transport.spi("0")?;
    let id = SpiFlash::read_jedec_id(&*spi, 16)?;
    log::info!("Read jedec id: {:x?}", id);
    let mut index = 0;
    while id[index] == 0x7f {
        index += 1;
    }
    let page = index;
    let manufacturer = id[index];
    let product = id[index + 1];
    let density = id[index + 2];

    assert_eq!(page, opts.jedec_page);
    assert_eq!(manufacturer, opts.jedec_id);
    assert_eq!(product, opts.jedec_product);
    assert_eq!(1u32 << density, opts.flash_size);
    Ok(())
}

fn test_write_enable_disable(transport: &TransportWrapper) -> Result<()> {
    let spi = transport.spi("0")?;

    assert_eq!(SpiFlash::read_status(&*spi)?, 0x0);

    SpiFlash::set_write_enable(&*spi)?;
    assert_eq!(SpiFlash::read_status(&*spi)?, 0x2);

    SpiFlash::set_write_disable(&*spi)?;
    assert_eq!(SpiFlash::read_status(&*spi)?, 0x0);

    Ok(())
}

fn test_sfdp(opts: &Opts, transport: &TransportWrapper) -> Result<()> {
    let spi = transport.spi("0")?;
    let sfdp = SpiFlash::read_sfdp(&*spi)?;
    log::info!("SFDP: {:#x?}", sfdp);

    // We expect a SFDP header of version 1.10 with no additional phdrs
    // except the mandatory JDEC phdr.
    assert_eq!(sfdp.header.signature, 0x50444653);
    assert_eq!(sfdp.header.minor, 0xa);
    assert_eq!(sfdp.header.major, 0x1);
    assert_eq!(sfdp.header.nph, 0x0);
    assert_eq!(sfdp.header.reserved, 0xff);

    // Check the parameters of the mandatory phdr.
    assert_eq!(sfdp.phdr[0].id, 0x0);
    assert_eq!(sfdp.phdr[0].minor, 0x7);
    assert_eq!(sfdp.phdr[0].major, 0x1);
    assert_eq!(sfdp.phdr[0].dwords, 0x17);
    assert_eq!(sfdp.phdr[0].offset, 0x10);

    // Check the values of the JEDEC parameters.
    // TODO(lowRISC/opentitan#14457): Update these checks after updating the
    // SFDP parser in opentitanlib.
    assert_eq!(sfdp.jedec.block_erase_size, BlockEraseSize::Block4KiB);
    assert_eq!(
        sfdp.jedec.write_granularity,
        WriteGranularity::Granularity64Bytes
    );
    assert_eq!(sfdp.jedec.write_en_required, true);
    assert_eq!(sfdp.jedec.write_en_opcode, 0x50);
    assert_eq!(sfdp.jedec.erase_opcode_4kib, 0x20);
    assert_eq!(sfdp.jedec.support_fast_read_112, false);
    assert_eq!(sfdp.jedec.address_modes, SupportedAddressModes::Mode3b);
    assert_eq!(sfdp.jedec.support_double_rate_clocking, false);
    assert_eq!(sfdp.jedec.support_fast_read_122, false);
    assert_eq!(sfdp.jedec.support_fast_read_144, false);
    assert_eq!(sfdp.jedec.support_fast_read_114, false);
    assert_eq!(sfdp.jedec.support_fast_read_222, false);
    assert_eq!(sfdp.jedec.support_fast_read_444, false);
    assert_eq!(sfdp.jedec.density, opts.flash_size);
    Ok(())
}

fn main() -> Result<()> {
    let opts = Opts::from_args();
    opts.init.init_logging();

    let transport = opts.init.init_target()?;
    execute_test!(test_bootstrap_enabled_not_requested, &opts, &transport);
    execute_test!(test_bootstrap_enabled_requested, &opts, &transport);
    execute_test!(test_jedec_id, &opts, &transport);
    execute_test!(test_sfdp, &opts, &transport);
    execute_test!(test_write_enable_disable, &transport);
    Ok(())
}
<commit_msg>[test] Use BootstrapTest to isolate individual test points<commit_after>// Copyright lowRISC contributors.
// Licensed under the Apache License, Version 2.0, see LICENSE for details.
// SPDX-License-Identifier: Apache-2.0

use anyhow::{bail, Result};
use regex::Regex;
use std::matches;
use std::ops::Drop;
use std::thread;
use std::time::Duration;
use structopt::StructOpt;

use opentitanlib::app::TransportWrapper;
use opentitanlib::execute_test;
use opentitanlib::spiflash::{
    sfdp, BlockEraseSize, SpiFlash, SupportedAddressModes, WriteGranularity,
};
use opentitanlib::test_utils::init::InitializeTest;
use opentitanlib::uart::console::{ExitStatus, UartConsole};
use opentitanlib::util::parse_int::ParseInt;

#[derive(Debug, StructOpt)]
struct Opts {
    #[structopt(flatten)]
    init: InitializeTest,

    #[structopt(
        long, parse(try_from_str=humantime::parse_duration),
        default_value = "5s",
        help = "Bootstrap detection timeout",
    )]
    timeout: Duration,

    #[structopt(
        long, parse(try_from_str=usize::from_str),
        default_value = "12",
        help = "JEDEC page of manufacturer",
    )]
    jedec_page: usize,

    #[structopt(
        long, parse(try_from_str=u8::from_str),
        default_value = "0xEF",
        help = "JEDEC ID of manufacturer",
    )]
    jedec_id: u8,

    #[structopt(
        long, parse(try_from_str=u8::from_str),
        default_value = "0x08",
        help = "JEDEC manufacturer product ID",
    )]
    jedec_product: u8,

    #[structopt(
        long, parse(try_from_str=u32::from_str),
        default_value = "0x100000",
        help = "Size of the internal flash",
    )]
    flash_size: u32,
}

/// A struct for isolating individual test points.
///
/// This is basically an example of the RAII technique.
/// Calling `BootstrapTest::start()` resets the chip and puts it in bootstrap mode (assuming it's
/// enabled). Strapping pins are released and the chip is reset when the returned struct is
/// dropped, typically at the end of the test point.
struct BootstrapTest<'a> {
    transport: &'a TransportWrapper,
    reset_delay: Duration,
}

impl<'a> BootstrapTest<'a> {
    pub fn start(transport: &'a TransportWrapper, reset_delay: Duration) -> Result<Self> {
        let b = BootstrapTest {
            transport,
            reset_delay,
        };
        transport.apply_pin_strapping("ROM_BOOTSTRAP")?;
        transport.reset_target(b.reset_delay, true)?;

        let spi = transport.spi("0")?;
        while SpiFlash::read_status(&*spi)? != 0 {
            thread::sleep(Duration::from_millis(10));
        }

        Ok(b)
    }
}

impl<'a> Drop for BootstrapTest<'a> {
    fn drop(&mut self) {
        self.transport
            .remove_pin_strapping("ROM_BOOTSTRAP")
            .unwrap();
        self.transport.reset_target(self.reset_delay, true).unwrap();
    }
}

fn test_bootstrap_enabled_requested(opts: &Opts, transport: &TransportWrapper) -> Result<()> {
    let _bs = BootstrapTest::start(transport, opts.init.bootstrap.options.reset_delay)?;

    // TODO: Should really `opts.init.uart_params.create()` here, but we need to refactor
    // BootstrapOptions first.
    //let uart = opts.init.uart_params.create(&transport)?;
    let uart = transport.uart("0")?;
    let mut console = UartConsole {
        timeout: Some(opts.timeout),
        exit_success: Some(Regex::new(r"bootstrap:1\r\n")?),
        exit_failure: Some(Regex::new(r"BFV:")?),
        ..Default::default()
    };

    transport.apply_pin_strapping("ROM_BOOTSTRAP")?;
    transport.reset_target(opts.init.bootstrap.options.reset_delay, true)?;

    // Now watch the console for the exit conditions.
    let result = console.interact(&*uart, None, Some(&mut std::io::stdout()))?;
    match result {
        ExitStatus::ExitSuccess => {}
        _ => {
            bail!("FAIL: {:?}", result);
        }
    };
    // Now check whether the SPI device is responding to status messages
    let spi = transport.spi("0")?;
    assert_eq!(SpiFlash::read_status(&*spi)?, 0x00);

    Ok(())
}

fn test_bootstrap_enabled_not_requested(opts: &Opts, transport: &TransportWrapper) -> Result<()> {
    // TODO: Should really `opts.init.uart_params.create()` here, but we need to refactor
    // BootstrapOptions first.
    //let uart = opts.init.uart_params.create(&transport)?;
    let uart = transport.uart("0")?;
    let mut console = UartConsole {
        timeout: Some(opts.timeout),
        exit_success: Some(Regex::new(r"BFV:")?),
        exit_failure: Some(Regex::new(r"bootstrap:1\r\n")?),
        ..Default::default()
    };

    transport.reset_target(opts.init.bootstrap.options.reset_delay, true)?;

    // Now watch the console for the exit conditions.
    let result = console.interact(&*uart, None, Some(&mut std::io::stdout()))?;
    match result {
        ExitStatus::ExitSuccess => {}
        _ => {
            bail!("FAIL: {:?}", result);
        }
    };

    // Now check whether the SPI device is responding to commands.
    // Note: CIPO line is in high-z state when CMD_INFO registers are not configured.
    // Use READ_SFDP instead of READ_STATUS to avoid false negatives when bootstrap is not
    // requested
    let spi = transport.spi("0")?;
    assert!(matches!(
        SpiFlash::read_sfdp(&*spi)
            .unwrap_err()
            .downcast::<sfdp::Error>()
            .unwrap(),
        sfdp::Error::WrongHeaderSignature(..)
    ));
    Ok(())
}

fn test_jedec_id(opts: &Opts, transport: &TransportWrapper) -> Result<()> {
    let _bs = BootstrapTest::start(transport, opts.init.bootstrap.options.reset_delay)?;

    let spi = transport.spi("0")?;
    let id = SpiFlash::read_jedec_id(&*spi, 16)?;
    log::info!("Read jedec id: {:x?}", id);
    let mut index = 0;
    while id[index] == 0x7f {
        index += 1;
    }
    let page = index;
    let manufacturer = id[index];
    let product = id[index + 1];
    let density = id[index + 2];

    assert_eq!(page, opts.jedec_page);
    assert_eq!(manufacturer, opts.jedec_id);
    assert_eq!(product, opts.jedec_product);
    assert_eq!(1u32 << density, opts.flash_size);
    Ok(())
}

fn test_write_enable_disable(opts: &Opts, transport: &TransportWrapper) -> Result<()> {
    let _bs = BootstrapTest::start(transport, opts.init.bootstrap.options.reset_delay)?;

    let spi = transport.spi("0")?;

    assert_eq!(SpiFlash::read_status(&*spi)?, 0x0);

    SpiFlash::set_write_enable(&*spi)?;
    assert_eq!(SpiFlash::read_status(&*spi)?, 0x2);

    SpiFlash::set_write_disable(&*spi)?;
    assert_eq!(SpiFlash::read_status(&*spi)?, 0x0);

    Ok(())
}

fn test_sfdp(opts: &Opts, transport: &TransportWrapper) -> Result<()> {
    let _bs = BootstrapTest::start(transport, opts.init.bootstrap.options.reset_delay)?;

    let spi = transport.spi("0")?;
    let sfdp = SpiFlash::read_sfdp(&*spi)?;
    log::info!("SFDP: {:#x?}", sfdp);

    // We expect a SFDP header of version 1.10 with no additional phdrs
    // except the mandatory JDEC phdr.
    assert_eq!(sfdp.header.signature, 0x50444653);
    assert_eq!(sfdp.header.minor, 0xa);
    assert_eq!(sfdp.header.major, 0x1);
    assert_eq!(sfdp.header.nph, 0x0);
    assert_eq!(sfdp.header.reserved, 0xff);

    // Check the parameters of the mandatory phdr.
    assert_eq!(sfdp.phdr[0].id, 0x0);
    assert_eq!(sfdp.phdr[0].minor, 0x7);
    assert_eq!(sfdp.phdr[0].major, 0x1);
    assert_eq!(sfdp.phdr[0].dwords, 0x17);
    assert_eq!(sfdp.phdr[0].offset, 0x10);

    // Check the values of the JEDEC parameters.
    // TODO(lowRISC/opentitan#14457): Update these checks after updating the
    // SFDP parser in opentitanlib.
    assert_eq!(sfdp.jedec.block_erase_size, BlockEraseSize::Block4KiB);
    assert_eq!(
        sfdp.jedec.write_granularity,
        WriteGranularity::Granularity64Bytes
    );
    assert_eq!(sfdp.jedec.write_en_required, true);
    assert_eq!(sfdp.jedec.write_en_opcode, 0x50);
    assert_eq!(sfdp.jedec.erase_opcode_4kib, 0x20);
    assert_eq!(sfdp.jedec.support_fast_read_112, false);
    assert_eq!(sfdp.jedec.address_modes, SupportedAddressModes::Mode3b);
    assert_eq!(sfdp.jedec.support_double_rate_clocking, false);
    assert_eq!(sfdp.jedec.support_fast_read_122, false);
    assert_eq!(sfdp.jedec.support_fast_read_144, false);
    assert_eq!(sfdp.jedec.support_fast_read_114, false);
    assert_eq!(sfdp.jedec.support_fast_read_222, false);
    assert_eq!(sfdp.jedec.support_fast_read_444, false);
    assert_eq!(sfdp.jedec.density, opts.flash_size);
    Ok(())
}

fn main() -> Result<()> {
    let opts = Opts::from_args();
    opts.init.init_logging();

    let transport = opts.init.init_target()?;
    execute_test!(test_bootstrap_enabled_not_requested, &opts, &transport);
    execute_test!(test_bootstrap_enabled_requested, &opts, &transport);
    execute_test!(test_jedec_id, &opts, &transport);
    execute_test!(test_sfdp, &opts, &transport);
    execute_test!(test_write_enable_disable, &opts, &transport);
    Ok(())
}
<|endoftext|>


<commit_before>use rendering::{Rectangle, Circle, Primitive, Renderable, Polygon, Line, LineShape};
use na::{Vector2, Vector3, Vector4, Rotation2, norm};
use geometry::{ConPoly, TwoDTransformable};
use geometry;

#[derive(Clone, Debug)]
pub struct Arrow {
    start: Vector2<f64>,
    end: Vector2<f64>,
    thickness: f64,
    arrow_dim: Vector2<f64>,
    color: Vector4<f64>,
    depth: f64,
    fixed: bool,
    shape: ArrowHeadShape
}

#[derive(Clone, Debug)]
pub enum ArrowHeadShape {
    Flat,
    RoundedLine
}

impl Arrow {
    pub fn new(
        start: Vector2<f64>,
        end: Vector2<f64>,
        thickness: f64,
        arrow_dim: Vector2<f64>,
        color: Vector4<f64>,
        depth: f64,
        fixed: bool
    ) -> Self {
        Arrow {
            start,
            end,
            thickness,
            arrow_dim,
            color,
            depth,
            fixed,
            shape: ArrowHeadShape::Flat
        }
    }

    pub fn new_rounded(
        start: Vector2<f64>,
        end: Vector2<f64>,
        thickness: f64,
        arrow_dim: Vector2<f64>,
        color: Vector4<f64>,
        depth: f64,
        fixed: bool
    ) -> Self {
        Arrow {
            start,
            end,
            thickness,
            arrow_dim,
            color,
            depth,
            fixed,
            shape: ArrowHeadShape::RoundedLine
        }
    }

    pub fn new_by_direction(
        start: Vector2<f64>,
        dir: Vector2<f64>,
        length: f64,
        thickness: f64,
        arrow_dim: Vector2<f64>,
        color: Vector4<f64>,
        depth: f64,
        fixed: bool
    ) -> Self {
        Arrow {
            start,
            end: start + length * dir,
            thickness,
            arrow_dim,
            color,
            depth,
            fixed,
            shape: ArrowHeadShape::Flat
        }
    }

    pub fn new_for_coll_test(
        start: Vector2<f64>,
        dir: Vector2<f64>,
        color: Vector4<f64>,
        depth: f64,
        fixed: bool
    ) -> Self {
        Arrow::new_by_direction(
            start,
            dir,
            0.04,
            0.01,
            Vector2::new(0.3, 0.02),
            color,
            depth,
            fixed
        )
    }

    fn get_center_line(&self) -> geometry::Line {
        geometry::Line::new(self.start, self.end)
    }

    fn get_length(&self) -> f64 {
        let shifted_end = self.end - self.start;
        shifted_end.norm()
    }

    fn get_line_angle(&self) -> f64 {
        let shifted_end = self.end - self.start;
        shifted_end.y.atan2(shifted_end.x)
    }

    fn generate_arrow_head(&self) -> Vec<Primitive> {
        match self.shape {
            ArrowHeadShape::Flat => self.generate_arrow_head_flat(),
            ArrowHeadShape::RoundedLine => self.generate_arrow_head_rounded()
        }
    }

    fn generate_arrow_head_flat(&self) -> Vec<Primitive> {
        let arrowhead_points = vec![
            Vector2::new(0.0, -self.arrow_dim.y),
            Vector2::new(self.arrow_dim.x, 0.0),
            Vector2::new(0.0, self.arrow_dim.y)
        ];
        let arrow_pos = self.get_center_line().get_point(1.0 - self.arrow_dim.x / self.get_length());

        let mut arrowhead = Polygon {
            corners: arrowhead_points,
            center: Vector2::new(0.0, 0.0),
            rot: Rotation2::new(self.get_line_angle()),
            pos: Vector3::new(arrow_pos.x, arrow_pos.y, self.depth),
            color: self.color,
            fixed: self.fixed
        };

        arrowhead.get_primitives()
    }

    pub fn generate_arrow_head_rounded(&self) -> Vec<Primitive> {
        let arrowhead_points = vec![
            Vector2::new(0.0, -self.arrow_dim.y),
            Vector2::new(self.arrow_dim.x, 0.0),
            Vector2::new(0.0, self.arrow_dim.y)
        ];
        let arrow_pos = self.get_center_line().get_point(1.0 - self.arrow_dim.x / self.get_length());
        let rotation = Rotation2::new(self.get_line_angle());

        let mut left_arrowhead_line = geometry::Line::new(
            rotation * arrowhead_points[0],
            rotation * arrowhead_points[1]);
        left_arrowhead_line.shift_by(arrow_pos);

        let mut left_line_renderable = Line::new_rounded(
            left_arrowhead_line.beg,
            left_arrowhead_line.end,
            self.thickness,
            self.color,
            self.depth,
            self.fixed);

        let mut right_arrowhead_line = geometry::Line::new(
            rotation * arrowhead_points[1],
            rotation * arrowhead_points[2]);
        right_arrowhead_line.shift_by(arrow_pos);

        let mut right_line_renderable = Line::new_rounded(
            right_arrowhead_line.beg,
            right_arrowhead_line.end,
            self.thickness,
            self.color,
            self.depth,
            self.fixed);

        let mut output = left_line_renderable.get_primitives();
        output.append(&mut right_line_renderable.get_primitives());
        output
    }
}

impl Renderable for Arrow {
    fn get_primitives(&mut self) -> Vec<Primitive> {
        let full_length = self.get_length();
        let center_line = geometry::Line::new(self.start, self.end);

        let root_line_length = match self.shape {
            ArrowHeadShape::Flat => full_length - self.arrow_dim.x,
            ArrowHeadShape::RoundedLine => full_length
        };

        let root_line_shape = match self.shape {
            ArrowHeadShape::Flat => LineShape::Square,
            ArrowHeadShape::RoundedLine => LineShape::Rounded
        };

        let mut root_line = Line::new(
            self.start,
            center_line.get_point(root_line_length / full_length),
            self.thickness,
            root_line_shape,
            self.color,
            self.depth,
            self.fixed
        );

        let mut output = root_line.get_primitives();
        output.append(&mut self.generate_arrow_head());
        output
    }
}<commit_msg>Minor change<commit_after>use rendering::{Rectangle, Circle, Primitive, Renderable, Polygon, Line, LineShape};
use na::{Vector2, Vector3, Vector4, Rotation2, norm};
use geometry::{ConPoly, TwoDTransformable};
use geometry;

#[derive(Clone, Debug)]
pub struct Arrow {
    pub start: Vector2<f64>,
    pub end: Vector2<f64>,
    thickness: f64,
    arrow_dim: Vector2<f64>,
    color: Vector4<f64>,
    depth: f64,
    fixed: bool,
    shape: ArrowHeadShape
}

#[derive(Clone, Debug)]
pub enum ArrowHeadShape {
    Flat,
    RoundedLine
}

impl Arrow {
    pub fn new(
        start: Vector2<f64>,
        end: Vector2<f64>,
        thickness: f64,
        arrow_dim: Vector2<f64>,
        color: Vector4<f64>,
        depth: f64,
        fixed: bool
    ) -> Self {
        Arrow {
            start,
            end,
            thickness,
            arrow_dim,
            color,
            depth,
            fixed,
            shape: ArrowHeadShape::Flat
        }
    }

    pub fn new_rounded(
        start: Vector2<f64>,
        end: Vector2<f64>,
        thickness: f64,
        arrow_dim: Vector2<f64>,
        color: Vector4<f64>,
        depth: f64,
        fixed: bool
    ) -> Self {
        Arrow {
            start,
            end,
            thickness,
            arrow_dim,
            color,
            depth,
            fixed,
            shape: ArrowHeadShape::RoundedLine
        }
    }

    pub fn new_by_direction(
        start: Vector2<f64>,
        dir: Vector2<f64>,
        length: f64,
        thickness: f64,
        arrow_dim: Vector2<f64>,
        color: Vector4<f64>,
        depth: f64,
        fixed: bool
    ) -> Self {
        Arrow {
            start,
            end: start + length * dir,
            thickness,
            arrow_dim,
            color,
            depth,
            fixed,
            shape: ArrowHeadShape::Flat
        }
    }

    pub fn new_for_coll_test(
        start: Vector2<f64>,
        dir: Vector2<f64>,
        color: Vector4<f64>,
        depth: f64,
        fixed: bool
    ) -> Self {
        Arrow::new_by_direction(
            start,
            dir,
            0.04,
            0.01,
            Vector2::new(0.3, 0.02),
            color,
            depth,
            fixed
        )
    }

    fn get_center_line(&self) -> geometry::Line {
        geometry::Line::new(self.start, self.end)
    }

    fn get_length(&self) -> f64 {
        let shifted_end = self.end - self.start;
        shifted_end.norm()
    }

    fn get_line_angle(&self) -> f64 {
        let shifted_end = self.end - self.start;
        shifted_end.y.atan2(shifted_end.x)
    }

    fn generate_arrow_head(&self) -> Vec<Primitive> {
        match self.shape {
            ArrowHeadShape::Flat => self.generate_arrow_head_flat(),
            ArrowHeadShape::RoundedLine => self.generate_arrow_head_rounded()
        }
    }

    fn generate_arrow_head_flat(&self) -> Vec<Primitive> {
        let arrowhead_points = vec![
            Vector2::new(0.0, -self.arrow_dim.y),
            Vector2::new(self.arrow_dim.x, 0.0),
            Vector2::new(0.0, self.arrow_dim.y)
        ];
        let arrow_pos = self.get_center_line().get_point(1.0 - self.arrow_dim.x / self.get_length());

        let mut arrowhead = Polygon {
            corners: arrowhead_points,
            center: Vector2::new(0.0, 0.0),
            rot: Rotation2::new(self.get_line_angle()),
            pos: Vector3::new(arrow_pos.x, arrow_pos.y, self.depth),
            color: self.color,
            fixed: self.fixed
        };

        arrowhead.get_primitives()
    }

    pub fn generate_arrow_head_rounded(&self) -> Vec<Primitive> {
        let arrowhead_points = vec![
            Vector2::new(0.0, -self.arrow_dim.y),
            Vector2::new(self.arrow_dim.x, 0.0),
            Vector2::new(0.0, self.arrow_dim.y)
        ];
        let arrow_pos = self.get_center_line().get_point(1.0 - self.arrow_dim.x / self.get_length());
        let rotation = Rotation2::new(self.get_line_angle());

        let mut left_arrowhead_line = geometry::Line::new(
            rotation * arrowhead_points[0],
            rotation * arrowhead_points[1]);
        left_arrowhead_line.shift_by(arrow_pos);

        let mut left_line_renderable = Line::new_rounded(
            left_arrowhead_line.beg,
            left_arrowhead_line.end,
            self.thickness,
            self.color,
            self.depth,
            self.fixed);

        let mut right_arrowhead_line = geometry::Line::new(
            rotation * arrowhead_points[1],
            rotation * arrowhead_points[2]);
        right_arrowhead_line.shift_by(arrow_pos);

        let mut right_line_renderable = Line::new_rounded(
            right_arrowhead_line.beg,
            right_arrowhead_line.end,
            self.thickness,
            self.color,
            self.depth,
            self.fixed);

        let mut output = left_line_renderable.get_primitives();
        output.append(&mut right_line_renderable.get_primitives());
        output
    }
}

impl Renderable for Arrow {
    fn get_primitives(&mut self) -> Vec<Primitive> {
        let full_length = self.get_length();
        let center_line = geometry::Line::new(self.start, self.end);

        let root_line_length = match self.shape {
            ArrowHeadShape::Flat => full_length - self.arrow_dim.x,
            ArrowHeadShape::RoundedLine => full_length
        };

        let root_line_shape = match self.shape {
            ArrowHeadShape::Flat => LineShape::Square,
            ArrowHeadShape::RoundedLine => LineShape::Rounded
        };

        let mut root_line = Line::new(
            self.start,
            center_line.get_point(root_line_length / full_length),
            self.thickness,
            root_line_shape,
            self.color,
            self.depth,
            self.fixed
        );

        let mut output = root_line.get_primitives();
        output.append(&mut self.generate_arrow_head());
        output
    }
}<|endoftext|>


<commit_before>/* Sample merge sort program in Rust.
   Tested to compile with rust-0.6-f1ddb8d.
*/
extern mod extra;
extern mod benchmark;
use std::vec;
use std::comm::{stream, Chan, Port};
use std::cell::Cell;
use benchmark::Benchmark;

fn merge_sort<T:Ord+Copy+Owned>(arr: ~[T]) -> ~[T] {
    let length = arr.len();
    if length <= 1 {
        return arr.to_owned();
    }

    let middle = length / 2;
    let mut left: ~[T] = vec::from_elem(middle, copy arr[0]);
    let mut right: ~[T] = vec::from_elem(length - middle, copy arr[0]);
    let mut index = 0;

    while index < middle {
        left[index] = arr[index];
        index += 1;
    }

    while index < length {
        right[index - middle] = arr[index];
        index += 1;
    }

    let (port, chan): (Port<~[T]>, Chan<~[T]>) = stream();
    let left_cell = Cell(left);
    do spawn {
        let sorted_left =  merge_sort(left_cell.take());
        chan.send(sorted_left);
    }
    right = merge_sort(right);

    left = port.recv();

    merge(left, right)
}

fn merge<T:Ord+Copy>(left_orig: ~[T], right_orig: ~[T]) -> ~[T] {
    let mut left = copy left_orig;
    let mut right = copy right_orig;
    let mut result = vec::from_elem(0, copy left[0]);

    while left.len() > 0 || right.len() > 0 {
        if left.len() > 0 && right.len() > 0 {
            if left[0] < right[0] {
                result.push(left.shift());
            }
            else {
                result.push(right.shift());
            }
        }
        else if left.len() > 0 {
            result.push(left.shift());
        }
        else {
            result.push(right.shift());
        }
    }

    return result;
}

fn main() {
    let mut bench = Benchmark::new();
    bench.run(merge_sort);
}

<commit_msg>Speed improvements to the merge-sort-par implementation.<commit_after>/* Sample merge sort program in Rust.
   Tested to compile with rust-0.6-f1ddb8d.
*/
extern mod extra;
extern mod benchmark;
use std::vec;
use std::comm::{stream, Chan, Port};
use std::cell::Cell;
use benchmark::Benchmark;

fn merge_sort<T:Ord+Copy+Owned>(arr: ~[T]) -> ~[T] {
    parallel_merge_sort(arr, 0)
}

fn parallel_merge_sort<T:Ord+Copy+Owned>(arr: ~[T], depth: uint) -> ~[T] {
    let length = arr.len();
    if length <= 1 {
        return arr.to_owned();
    }

    let middle = length / 2;
    let mut left: ~[T] = vec::from_elem(middle, copy arr[0]);
    let mut right: ~[T] = vec::from_elem(length - middle, copy arr[0]);
    let mut index = 0;

    while index < middle {
        left[index] = arr[index];
        index += 1;
    }

    while index < length {
        right[index - middle] = arr[index];
        index += 1;
    }

    if depth < 8 {
        let (port, chan): (Port<~[T]>, Chan<~[T]>) = stream();
        let left_cell = Cell(left);
        do spawn {
            let sorted_left =  parallel_merge_sort(left_cell.take(), depth + 1);
            chan.send(sorted_left);
        }
        right = parallel_merge_sort(right, depth + 1);

        left = port.recv();
    } else {
        left = parallel_merge_sort(left, depth);
        right = parallel_merge_sort(right, depth);
    }

    merge(left, right)
}

fn merge<T:Ord+Copy>(left_orig: ~[T], right_orig: ~[T]) -> ~[T] {
    let mut left = copy left_orig;
    let mut right = copy right_orig;
    let mut result = vec::from_elem(0, copy left[0]);

    while left.len() > 0 || right.len() > 0 {
        if left.len() > 0 && right.len() > 0 {
            if left[0] < right[0] {
                result.push(left.shift());
            }
            else {
                result.push(right.shift());
            }
        }
        else if left.len() > 0 {
            result.push(left.shift());
        }
        else {
            result.push(right.shift());
        }
    }

    return result;
}

fn main() {
    let mut bench = Benchmark::new();
    bench.run(merge_sort);
}

<|endoftext|>


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

use aptos_rest_client::Client as RestClient;
use async_trait::async_trait;
use core::time;
use futures::future::join_all;
use std::time::Duration;
use std::{collections::HashSet, time::Instant};

use crate::{wait_for_all_nodes_to_catchup_to_version, Swarm, SwarmExt};

#[derive(Clone, Debug)]
pub struct NodeState {
    pub version: u64,
    pub epoch: u64,
    pub round: u64,
}

// TODO: check if we can fetch consensus round, not just committed round.
async fn get_node_state(validator_client: &RestClient) -> NodeState {
    let (events, state) = validator_client
        .get_new_block_events(None, Some(1))
        .await
        .unwrap()
        .into_parts();
    let event = events.first().unwrap();
    assert!(event.version <= state.version);
    NodeState {
        version: state.version,
        epoch: event.event.epoch(),
        round: event.event.round(),
    }
}

/// Run a test, where we spin up a set of validators, and then
/// check a reliability scenario.
/// After scenario finishes, we always return reliability to 100%,
/// and we confirm that chain recouperates, makes progress, and
/// all nodes agree on ledger.
///
/// Scenario is performed via two nested loops:
/// outer cycles, and within each cycle over parts.
/// Scenario can specify failpoint changes on every part,
/// and can check the performance of the network after every cycle.
///
/// Transaction can be inserted on every part, to control the throughput.
/// I.e. if part is shorter than how long it takes for empty block to be
/// generated, we can make sure one block gets created on every part.
pub async fn test_consensus_fault_tolerance(
    swarm: &mut dyn Swarm,
    cycles: usize,
    cycle_duration_s: f32,
    parts_in_cycle: usize,
    mut failure_injection: Box<dyn FailureInjection>,
    // (cycle, executed_epochs, executed_rounds, executed_transactions, current_state, previous_state)
    mut check_cycle: Box<dyn FnMut(usize, u64, u64, u64, Vec<NodeState>, Vec<NodeState>)>,
    new_epoch_on_cycle: bool,
) -> anyhow::Result<()> {
    let validator_clients = swarm.get_clients_with_names();

    async fn get_all_states(validator_clients: &[(String, RestClient)]) -> Vec<NodeState> {
        join_all(
            validator_clients
                .iter()
                .cloned()
                .map(move |(_, v)| async move { get_node_state(&v).await }),
        )
        .await
    }

    for cycle in 0..cycles {
        let previous = get_all_states(&validator_clients).await;

        let now = Instant::now();
        for part in 0..parts_in_cycle {
            failure_injection
                .inject(&validator_clients, cycle, part)
                .await;
            let elapsed = now.elapsed().as_secs_f32();
            let wanted = (1 + part) as f32 * cycle_duration_s / (parts_in_cycle as f32);
            if elapsed < wanted {
                tokio::time::sleep(time::Duration::from_secs_f32(wanted - elapsed)).await;
            }
        }

        let cur = get_all_states(&validator_clients).await;

        let epochs = cur.iter().map(|s| s.epoch).max().unwrap()
            - previous.iter().map(|s| s.epoch).max().unwrap();
        let rounds = cur
            .iter()
            .map(|s| s.round)
            .max()
            .unwrap()
            .saturating_sub(previous.iter().map(|s| s.round).max().unwrap());
        let transactions = cur.iter().map(|s| s.version).max().unwrap()
            - previous.iter().map(|s| s.version).max().unwrap();

        println!(
            "cycle {} lasted {:.3} with {} epochs, {} rounds and {} transactions",
            cycle,
            now.elapsed().as_secs_f32(),
            epochs,
            rounds,
            transactions,
        );
        println!(
            "All at versions: {:?}",
            cur.iter().map(|s| s.version).collect::<Vec<_>>()
        );
        println!(
            "All at rounds: {:?}",
            cur.iter().map(|s| s.round).collect::<Vec<_>>()
        );

        check_cycle(cycle, epochs, rounds, transactions, cur.clone(), previous);
        if new_epoch_on_cycle {
            swarm.aptos_public_info().reconfig().await;
        }
    }

    failure_injection.clear(&validator_clients).await;

    let cur = get_all_states(&validator_clients).await;
    println!(
        "All at versions: {:?}",
        cur.iter().map(|s| s.version).collect::<Vec<_>>()
    );
    let largest_v = cur.iter().map(|s| s.version).max().unwrap();
    println!("Largest version {}", largest_v);
    let target_v = largest_v + 10;

    wait_for_all_nodes_to_catchup_to_version(&validator_clients, target_v, Duration::from_secs(30))
        .await?;

    let transactions: Vec<_> =
        join_all(validator_clients.iter().cloned().map(move |v| async move {
            let mut txns =
                v.1.get_transactions(None, Some(1000))
                    .await
                    .map_err(|e| anyhow::anyhow!("{:?}", e))?
                    .into_inner();
            txns.retain(|t| t.version().unwrap() <= target_v);
            <anyhow::Result<Vec<_>>>::Ok(txns)
        }))
        .await;

    for i in 1..transactions.len() {
        let txns_a = transactions
            .first()
            .unwrap()
            .as_ref()
            .map_err(|e| anyhow::anyhow!("{:?}", &e))?;
        let txns_b = transactions
            .get(i)
            .unwrap()
            .as_ref()
            .map_err(|e| anyhow::anyhow!("{:?}", &e))?;
        assert_eq!(txns_a.len(), txns_b.len());
        for i in 0..txns_a.len() {
            assert_eq!(txns_a[i], txns_b[i]);
        }
    }

    Ok(())
}

#[async_trait]
pub trait FailureInjection {
    async fn inject(
        &mut self,
        validator_clients: &[(String, RestClient)],
        cycle: usize,
        part: usize,
    );
    async fn clear(&mut self, validator_clients: &[(String, RestClient)]);
}

pub struct NoFailureInjection {}

#[async_trait]
impl FailureInjection for NoFailureInjection {
    async fn inject(&mut self, _: &[(String, RestClient)], _: usize, _: usize) {}
    async fn clear(&mut self, _: &[(String, RestClient)]) {}
}

pub fn no_failure_injection() -> Box<dyn FailureInjection> {
    Box::new(NoFailureInjection {})
}

pub struct FailPointFailureInjection {
    modified_failpoints: HashSet<(usize, String)>,
    // (cycle, part) -> (Vec(validator_index, name, action), reset_old_enpoints)
    get_fail_points_to_set:
        Box<dyn FnMut(usize, usize) -> (Vec<(usize, String, String)>, bool) + Send>,
}

impl FailPointFailureInjection {
    pub fn new(
        get_fail_points_to_set: Box<
            dyn FnMut(usize, usize) -> (Vec<(usize, String, String)>, bool) + Send,
        >,
    ) -> Self {
        Self {
            modified_failpoints: HashSet::new(),
            get_fail_points_to_set,
        }
    }
}

pub fn fail_point_injection(
    get_fail_points_to_set: Box<
        dyn FnMut(usize, usize) -> (Vec<(usize, String, String)>, bool) + Send,
    >,
) -> Box<dyn FailureInjection> {
    Box::new(FailPointFailureInjection::new(get_fail_points_to_set))
}

#[async_trait]
impl FailureInjection for FailPointFailureInjection {
    async fn inject(
        &mut self,
        validator_clients: &[(String, RestClient)],
        cycle: usize,
        part: usize,
    ) {
        let (fail_points_to_set, reset_old_failpoints) = (self.get_fail_points_to_set)(cycle, part);
        if reset_old_failpoints {
            for (validator_idx, name) in self.modified_failpoints.iter() {
                validator_clients[*validator_idx]
                    .1
                    .set_failpoint(name.clone(), "off".to_string())
                    .await
                    .unwrap();
            }
            self.modified_failpoints = HashSet::new();
        }
        for (validator_idx, name, actions) in fail_points_to_set {
            validator_clients[validator_idx]
                .1
                .set_failpoint(name.clone(), actions.clone())
                .await
                .unwrap();
            self.modified_failpoints.insert((validator_idx, name));
        }
    }
    async fn clear(&mut self, validator_clients: &[(String, RestClient)]) {
        for (validator_idx, name) in self.modified_failpoints.iter() {
            validator_clients[*validator_idx]
                .1
                .set_failpoint(name.clone(), "off".to_string())
                .await
                .unwrap();
        }
    }
}
<commit_msg>[forge] fix consensus stress test check (#3902)<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use aptos_rest_client::Client as RestClient;
use async_trait::async_trait;
use core::time;
use futures::future::join_all;
use std::time::Duration;
use std::{collections::HashSet, time::Instant};

use crate::{wait_for_all_nodes_to_catchup_to_version, Swarm, SwarmExt};

#[derive(Clone, Debug)]
pub struct NodeState {
    pub version: u64,
    pub epoch: u64,
    pub round: u64,
}

// TODO: check if we can fetch consensus round, not just committed round.
async fn get_node_state(validator_client: &RestClient) -> NodeState {
    let (events, state) = validator_client
        .get_new_block_events(None, Some(1))
        .await
        .unwrap()
        .into_parts();
    let event = events.first().unwrap();
    assert!(event.version <= state.version);
    NodeState {
        version: state.version,
        epoch: event.event.epoch(),
        round: event.event.round(),
    }
}

/// Run a test, where we spin up a set of validators, and then
/// check a reliability scenario.
/// After scenario finishes, we always return reliability to 100%,
/// and we confirm that chain recouperates, makes progress, and
/// all nodes agree on ledger.
///
/// Scenario is performed via two nested loops:
/// outer cycles, and within each cycle over parts.
/// Scenario can specify failpoint changes on every part,
/// and can check the performance of the network after every cycle.
///
/// Transaction can be inserted on every part, to control the throughput.
/// I.e. if part is shorter than how long it takes for empty block to be
/// generated, we can make sure one block gets created on every part.
pub async fn test_consensus_fault_tolerance(
    swarm: &mut dyn Swarm,
    cycles: usize,
    cycle_duration_s: f32,
    parts_in_cycle: usize,
    mut failure_injection: Box<dyn FailureInjection>,
    // (cycle, executed_epochs, executed_rounds, executed_transactions, current_state, previous_state)
    mut check_cycle: Box<dyn FnMut(usize, u64, u64, u64, Vec<NodeState>, Vec<NodeState>)>,
    new_epoch_on_cycle: bool,
) -> anyhow::Result<()> {
    let validator_clients = swarm.get_clients_with_names();

    async fn get_all_states(validator_clients: &[(String, RestClient)]) -> Vec<NodeState> {
        join_all(
            validator_clients
                .iter()
                .cloned()
                .map(move |(_, v)| async move { get_node_state(&v).await }),
        )
        .await
    }

    for cycle in 0..cycles {
        let previous = get_all_states(&validator_clients).await;

        let now = Instant::now();
        for part in 0..parts_in_cycle {
            failure_injection
                .inject(&validator_clients, cycle, part)
                .await;
            let elapsed = now.elapsed().as_secs_f32();
            let wanted = (1 + part) as f32 * cycle_duration_s / (parts_in_cycle as f32);
            if elapsed < wanted {
                tokio::time::sleep(time::Duration::from_secs_f32(wanted - elapsed)).await;
            }
        }

        let cur = get_all_states(&validator_clients).await;

        let epochs = cur.iter().map(|s| s.epoch).max().unwrap()
            - previous.iter().map(|s| s.epoch).max().unwrap();
        let rounds = cur
            .iter()
            .map(|s| s.round)
            .max()
            .unwrap()
            .saturating_sub(previous.iter().map(|s| s.round).max().unwrap());
        let transactions = cur.iter().map(|s| s.version).max().unwrap()
            - previous.iter().map(|s| s.version).max().unwrap();

        println!(
            "cycle {} lasted {:.3} with {} epochs, {} rounds and {} transactions",
            cycle,
            now.elapsed().as_secs_f32(),
            epochs,
            rounds,
            transactions,
        );
        println!(
            "All at versions: {:?}",
            cur.iter().map(|s| s.version).collect::<Vec<_>>()
        );
        println!(
            "All at rounds: {:?}",
            cur.iter().map(|s| s.round).collect::<Vec<_>>()
        );

        check_cycle(cycle, epochs, rounds, transactions, cur.clone(), previous);
        if new_epoch_on_cycle {
            swarm.aptos_public_info().reconfig().await;
        }
    }

    failure_injection.clear(&validator_clients).await;

    let cur = get_all_states(&validator_clients).await;
    println!(
        "All at versions: {:?}",
        cur.iter().map(|s| s.version).collect::<Vec<_>>()
    );
    let largest_v = cur.iter().map(|s| s.version).max().unwrap();
    println!("Largest version {}", largest_v);
    let target_v = largest_v + 10;

    wait_for_all_nodes_to_catchup_to_version(&validator_clients, target_v, Duration::from_secs(30))
        .await?;

    let transactions: Vec<_> =
        join_all(validator_clients.iter().cloned().map(move |v| async move {
            let mut txns =
                v.1.get_transactions_bcs(Some(target_v.saturating_sub(1000)), Some(1000))
                    .await
                    .map_err(|e| anyhow::anyhow!("{:?}", e))?
                    .into_inner();
            txns.retain(|t| t.version <= target_v);
            <anyhow::Result<Vec<_>>>::Ok(txns)
        }))
        .await;

    let txns_a = transactions
        .first()
        .unwrap()
        .as_ref()
        .map_err(|e| anyhow::anyhow!("{:?}", &e))?;

    for i in 1..transactions.len() {
        let txns_b = transactions
            .get(i)
            .unwrap()
            .as_ref()
            .map_err(|e| anyhow::anyhow!("{:?}", &e))?;
        assert_eq!(
            txns_a.len(),
            txns_b.len(),
            "Fetched length of transactions for target_v {} doesn't match: from {} to {} vs from {} to {}",
            target_v,
            txns_a.first().map(|t| t.version).unwrap_or(0),
            txns_a.last().map(|t| t.version).unwrap_or(0),
            txns_b.first().map(|t| t.version).unwrap_or(0),
            txns_b.last().map(|t| t.version).unwrap_or(0),
        );
        for i in 0..txns_a.len() {
            assert_eq!(
                txns_a[i], txns_b[i],
                "Transaction at index {} after target version {}, doesn't match",
                i, target_v
            );
        }
    }

    Ok(())
}

#[async_trait]
pub trait FailureInjection {
    async fn inject(
        &mut self,
        validator_clients: &[(String, RestClient)],
        cycle: usize,
        part: usize,
    );
    async fn clear(&mut self, validator_clients: &[(String, RestClient)]);
}

pub struct NoFailureInjection {}

#[async_trait]
impl FailureInjection for NoFailureInjection {
    async fn inject(&mut self, _: &[(String, RestClient)], _: usize, _: usize) {}
    async fn clear(&mut self, _: &[(String, RestClient)]) {}
}

pub fn no_failure_injection() -> Box<dyn FailureInjection> {
    Box::new(NoFailureInjection {})
}

pub struct FailPointFailureInjection {
    modified_failpoints: HashSet<(usize, String)>,
    // (cycle, part) -> (Vec(validator_index, name, action), reset_old_enpoints)
    get_fail_points_to_set:
        Box<dyn FnMut(usize, usize) -> (Vec<(usize, String, String)>, bool) + Send>,
}

impl FailPointFailureInjection {
    pub fn new(
        get_fail_points_to_set: Box<
            dyn FnMut(usize, usize) -> (Vec<(usize, String, String)>, bool) + Send,
        >,
    ) -> Self {
        Self {
            modified_failpoints: HashSet::new(),
            get_fail_points_to_set,
        }
    }
}

pub fn fail_point_injection(
    get_fail_points_to_set: Box<
        dyn FnMut(usize, usize) -> (Vec<(usize, String, String)>, bool) + Send,
    >,
) -> Box<dyn FailureInjection> {
    Box::new(FailPointFailureInjection::new(get_fail_points_to_set))
}

#[async_trait]
impl FailureInjection for FailPointFailureInjection {
    async fn inject(
        &mut self,
        validator_clients: &[(String, RestClient)],
        cycle: usize,
        part: usize,
    ) {
        let (fail_points_to_set, reset_old_failpoints) = (self.get_fail_points_to_set)(cycle, part);
        if reset_old_failpoints {
            for (validator_idx, name) in self.modified_failpoints.iter() {
                validator_clients[*validator_idx]
                    .1
                    .set_failpoint(name.clone(), "off".to_string())
                    .await
                    .unwrap();
            }
            self.modified_failpoints = HashSet::new();
        }
        for (validator_idx, name, actions) in fail_points_to_set {
            validator_clients[validator_idx]
                .1
                .set_failpoint(name.clone(), actions.clone())
                .await
                .unwrap();
            self.modified_failpoints.insert((validator_idx, name));
        }
    }
    async fn clear(&mut self, validator_clients: &[(String, RestClient)]) {
        for (validator_idx, name) in self.modified_failpoints.iter() {
            validator_clients[*validator_idx]
                .1
                .set_failpoint(name.clone(), "off".to_string())
                .await
                .unwrap();
        }
    }
}
<|endoftext|>


<commit_before>//! Handles the `Enter` key press. At the momently, this only continues
//! comments, but should handle indent some time in the future as well.

use ide_db::base_db::{FilePosition, SourceDatabase};
use ide_db::RootDatabase;
use syntax::{
    algo::find_node_at_offset,
    ast::{self, edit::IndentLevel, AstToken},
    AstNode, SmolStr, SourceFile,
    SyntaxKind::*,
    SyntaxNode, SyntaxToken, TextRange, TextSize, TokenAtOffset,
};

use text_edit::TextEdit;

// Feature: On Enter
//
// rust-analyzer can override kbd:[Enter] key to make it smarter:
//
// - kbd:[Enter] inside triple-slash comments automatically inserts `///`
// - kbd:[Enter] in the middle or after a trailing space in `//` inserts `//`
// - kbd:[Enter] inside `//!` doc comments automatically inserts `//!`
// - kbd:[Enter] after `{` indents contents and closing `}` of single-line block
//
// This action needs to be assigned to shortcut explicitly.
//
// VS Code::
//
// Add the following to `keybindings.json`:
// [source,json]
// ----
// {
//   "key": "Enter",
//   "command": "rust-analyzer.onEnter",
//   "when": "editorTextFocus && !suggestWidgetVisible && editorLangId == rust"
// }
// ----
//
// When using the Vim plugin:
// [source,json]
// ----
// {
//   "key": "Enter",
//   "command": "rust-analyzer.onEnter",
//   "when": "editorTextFocus && !suggestWidgetVisible && editorLangId == rust && vim.mode == 'Insert'"
// }
// ----
//
// image::https://user-images.githubusercontent.com/48062697/113065578-04c21800-91b1-11eb-82b8-22b8c481e645.gif[]
pub(crate) fn on_enter(db: &RootDatabase, position: FilePosition) -> Option<TextEdit> {
    let parse = db.parse(position.file_id);
    let file = parse.tree();
    let token = file.syntax().token_at_offset(position.offset).left_biased()?;

    if let Some(comment) = ast::Comment::cast(token.clone()) {
        return on_enter_in_comment(&comment, &file, position.offset);
    }

    if token.kind() == L_CURLY {
        // Typing enter after the `{` of a block expression, where the `}` is on the same line
        if let Some(edit) = find_node_at_offset(file.syntax(), position.offset - TextSize::of('{'))
            .and_then(|block| on_enter_in_block(block, position))
        {
            cov_mark::hit!(indent_block_contents);
            return Some(edit);
        }

        // Typing enter after the `{` of a use tree list.
        if let Some(edit) = find_node_at_offset(file.syntax(), position.offset - TextSize::of('{'))
            .and_then(|list| on_enter_in_use_tree_list(list, position))
        {
            cov_mark::hit!(indent_block_contents);
            return Some(edit);
        }
    }

    None
}

fn on_enter_in_comment(
    comment: &ast::Comment,
    file: &ast::SourceFile,
    offset: TextSize,
) -> Option<TextEdit> {
    if comment.kind().shape.is_block() {
        return None;
    }

    let prefix = comment.prefix();
    let comment_range = comment.syntax().text_range();
    if offset < comment_range.start() + TextSize::of(prefix) {
        return None;
    }

    let mut remove_trailing_whitespace = false;
    // Continuing single-line non-doc comments (like this one :) ) is annoying
    if prefix == "//" && comment_range.end() == offset {
        if comment.text().ends_with(' ') {
            cov_mark::hit!(continues_end_of_line_comment_with_space);
            remove_trailing_whitespace = true;
        } else if !followed_by_comment(comment) {
            return None;
        }
    }

    let indent = node_indent(file, comment.syntax())?;
    let inserted = format!("\n{}{} $0", indent, prefix);
    let delete = if remove_trailing_whitespace {
        let trimmed_len = comment.text().trim_end().len() as u32;
        let trailing_whitespace_len = comment.text().len() as u32 - trimmed_len;
        TextRange::new(offset - TextSize::from(trailing_whitespace_len), offset)
    } else {
        TextRange::empty(offset)
    };
    let edit = TextEdit::replace(delete, inserted);
    Some(edit)
}

fn on_enter_in_block(block: ast::BlockExpr, position: FilePosition) -> Option<TextEdit> {
    let contents = block_contents(&block)?;

    if block.syntax().text().contains_char('\n') {
        return None;
    }

    let indent = IndentLevel::from_node(block.syntax());
    let mut edit = TextEdit::insert(position.offset, format!("\n{}$0", indent + 1));
    edit.union(TextEdit::insert(contents.text_range().end(), format!("\n{}", indent))).ok()?;
    Some(edit)
}

fn on_enter_in_use_tree_list(list: ast::UseTreeList, position: FilePosition) -> Option<TextEdit> {
    if list.syntax().text().contains_char('\n') {
        return None;
    }

    let indent = IndentLevel::from_node(list.syntax());
    let mut edit = TextEdit::insert(position.offset, format!("\n{}$0", indent + 1));
    edit.union(TextEdit::insert(
        list.r_curly_token()?.text_range().start(),
        format!("\n{}", indent),
    ))
    .ok()?;
    Some(edit)
}

fn block_contents(block: &ast::BlockExpr) -> Option<SyntaxNode> {
    let mut node = block.tail_expr().map(|e| e.syntax().clone());

    for stmt in block.statements() {
        if node.is_some() {
            // More than 1 node in the block
            return None;
        }

        node = Some(stmt.syntax().clone());
    }

    node
}

fn followed_by_comment(comment: &ast::Comment) -> bool {
    let ws = match comment.syntax().next_token().and_then(ast::Whitespace::cast) {
        Some(it) => it,
        None => return false,
    };
    if ws.spans_multiple_lines() {
        return false;
    }
    ws.syntax().next_token().and_then(ast::Comment::cast).is_some()
}

fn node_indent(file: &SourceFile, token: &SyntaxToken) -> Option<SmolStr> {
    let ws = match file.syntax().token_at_offset(token.text_range().start()) {
        TokenAtOffset::Between(l, r) => {
            assert!(r == *token);
            l
        }
        TokenAtOffset::Single(n) => {
            assert!(n == *token);
            return Some("".into());
        }
        TokenAtOffset::None => unreachable!(),
    };
    if ws.kind() != WHITESPACE {
        return None;
    }
    let text = ws.text();
    let pos = text.rfind('\n').map(|it| it + 1).unwrap_or(0);
    Some(text[pos..].into())
}

#[cfg(test)]
mod tests {
    use stdx::trim_indent;
    use test_utils::assert_eq_text;

    use crate::fixture;

    fn apply_on_enter(before: &str) -> Option<String> {
        let (analysis, position) = fixture::position(before);
        let result = analysis.on_enter(position).unwrap()?;

        let mut actual = analysis.file_text(position.file_id).unwrap().to_string();
        result.apply(&mut actual);
        Some(actual)
    }

    fn do_check(ra_fixture_before: &str, ra_fixture_after: &str) {
        let ra_fixture_after = &trim_indent(ra_fixture_after);
        let actual = apply_on_enter(ra_fixture_before).unwrap();
        assert_eq_text!(ra_fixture_after, &actual);
    }

    fn do_check_noop(ra_fixture_text: &str) {
        assert!(apply_on_enter(ra_fixture_text).is_none())
    }

    #[test]
    fn continues_doc_comment() {
        do_check(
            r"
/// Some docs$0
fn foo() {
}
",
            r"
/// Some docs
/// $0
fn foo() {
}
",
        );

        do_check(
            r"
impl S {
    /// Some$0 docs.
    fn foo() {}
}
",
            r"
impl S {
    /// Some
    /// $0 docs.
    fn foo() {}
}
",
        );

        do_check(
            r"
///$0 Some docs
fn foo() {
}
",
            r"
///
/// $0 Some docs
fn foo() {
}
",
        );
    }

    #[test]
    fn does_not_continue_before_doc_comment() {
        do_check_noop(r"$0//! docz");
    }

    #[test]
    fn continues_another_doc_comment() {
        do_check(
            r#"
fn main() {
    //! Documentation for$0 on enter
    let x = 1 + 1;
}
"#,
            r#"
fn main() {
    //! Documentation for
    //! $0 on enter
    let x = 1 + 1;
}
"#,
        );
    }

    #[test]
    fn continues_code_comment_in_the_middle_of_line() {
        do_check(
            r"
fn main() {
    // Fix$0 me
    let x = 1 + 1;
}
",
            r"
fn main() {
    // Fix
    // $0 me
    let x = 1 + 1;
}
",
        );
    }

    #[test]
    fn continues_code_comment_in_the_middle_several_lines() {
        do_check(
            r"
fn main() {
    // Fix$0
    // me
    let x = 1 + 1;
}
",
            r"
fn main() {
    // Fix
    // $0
    // me
    let x = 1 + 1;
}
",
        );
    }

    #[test]
    fn does_not_continue_end_of_line_comment() {
        do_check_noop(
            r"
fn main() {
    // Fix me$0
    let x = 1 + 1;
}
",
        );
    }

    #[test]
    fn continues_end_of_line_comment_with_space() {
        cov_mark::check!(continues_end_of_line_comment_with_space);
        do_check(
            r#"
fn main() {
    // Fix me $0
    let x = 1 + 1;
}
"#,
            r#"
fn main() {
    // Fix me
    // $0
    let x = 1 + 1;
}
"#,
        );
    }

    #[test]
    fn trims_all_trailing_whitespace() {
        do_check(
            "
fn main() {
    // Fix me  \t\t   $0
    let x = 1 + 1;
}
",
            "
fn main() {
    // Fix me
    // $0
    let x = 1 + 1;
}
",
        );
    }

    #[test]
    fn indents_fn_body_block() {
        cov_mark::check!(indent_block_contents);
        do_check(
            r#"
fn f() {$0()}
        "#,
            r#"
fn f() {
    $0()
}
        "#,
        );
    }

    #[test]
    fn indents_block_expr() {
        do_check(
            r#"
fn f() {
    let x = {$0()};
}
        "#,
            r#"
fn f() {
    let x = {
        $0()
    };
}
        "#,
        );
    }

    #[test]
    fn indents_match_arm() {
        do_check(
            r#"
fn f() {
    match 6 {
        1 => {$0f()},
        _ => (),
    }
}
        "#,
            r#"
fn f() {
    match 6 {
        1 => {
            $0f()
        },
        _ => (),
    }
}
        "#,
        );
    }

    #[test]
    fn indents_block_with_statement() {
        do_check(
            r#"
fn f() {$0a = b}
        "#,
            r#"
fn f() {
    $0a = b
}
        "#,
        );
        do_check(
            r#"
fn f() {$0fn f() {}}
        "#,
            r#"
fn f() {
    $0fn f() {}
}
        "#,
        );
    }

    #[test]
    fn indents_nested_blocks() {
        do_check(
            r#"
fn f() {$0{}}
        "#,
            r#"
fn f() {
    $0{}
}
        "#,
        );
    }

    #[test]
    fn does_not_indent_empty_block() {
        do_check_noop(
            r#"
fn f() {$0}
        "#,
        );
        do_check_noop(
            r#"
fn f() {{$0}}
        "#,
        );
    }

    #[test]
    fn does_not_indent_block_with_too_much_content() {
        do_check_noop(
            r#"
fn f() {$0 a = b; ()}
        "#,
        );
        do_check_noop(
            r#"
fn f() {$0 a = b; a = b; }
        "#,
        );
    }

    #[test]
    fn does_not_indent_multiline_block() {
        do_check_noop(
            r#"
fn f() {$0
}
        "#,
        );
        do_check_noop(
            r#"
fn f() {$0

}
        "#,
        );
    }

    #[test]
    fn indents_use_tree_list() {
        do_check(
            r#"
use crate::{$0};
            "#,
            r#"
use crate::{
    $0
};
            "#,
        );
        do_check(
            r#"
use crate::{$0Object, path::to::OtherThing};
            "#,
            r#"
use crate::{
    $0Object, path::to::OtherThing
};
            "#,
        );
        do_check(
            r#"
use {crate::{$0Object, path::to::OtherThing}};
            "#,
            r#"
use {crate::{
    $0Object, path::to::OtherThing
}};
            "#,
        );
        do_check(
            r#"
use {
    crate::{$0Object, path::to::OtherThing}
};
            "#,
            r#"
use {
    crate::{
        $0Object, path::to::OtherThing
    }
};
            "#,
        );
    }

    #[test]
    fn does_not_indent_use_tree_list_when_not_at_curly_brace() {
        do_check_noop(
            r#"
use path::{Thing$0};
            "#,
        );
    }

    #[test]
    fn does_not_indent_use_tree_list_without_curly_braces() {
        do_check_noop(
            r#"
use path::Thing$0;
            "#,
        );
        do_check_noop(
            r#"
use path::$0Thing;
            "#,
        );
        do_check_noop(
            r#"
use path::Thing$0};
            "#,
        );
        do_check_noop(
            r#"
use path::{$0Thing;
            "#,
        );
    }

    #[test]
    fn does_not_indent_multiline_use_tree_list() {
        do_check_noop(
            r#"
use path::{$0
    Thing
};
            "#,
        );
    }
}
<commit_msg>Add scary warning to onEnter<commit_after>//! Handles the `Enter` key press. At the momently, this only continues
//! comments, but should handle indent some time in the future as well.

use ide_db::base_db::{FilePosition, SourceDatabase};
use ide_db::RootDatabase;
use syntax::{
    algo::find_node_at_offset,
    ast::{self, edit::IndentLevel, AstToken},
    AstNode, SmolStr, SourceFile,
    SyntaxKind::*,
    SyntaxNode, SyntaxToken, TextRange, TextSize, TokenAtOffset,
};

use text_edit::TextEdit;

// Feature: On Enter
//
// rust-analyzer can override kbd:[Enter] key to make it smarter:
//
// - kbd:[Enter] inside triple-slash comments automatically inserts `///`
// - kbd:[Enter] in the middle or after a trailing space in `//` inserts `//`
// - kbd:[Enter] inside `//!` doc comments automatically inserts `//!`
// - kbd:[Enter] after `{` indents contents and closing `}` of single-line block
//
// This action needs to be assigned to shortcut explicitly.
//
// Note that, depending on the other installed extensions, this feature can visibly slow down typing.
// Similarly, if rust-analyzer crashes or stops responding, `Enter` might not work.
// In that case, you can still press `Shift-Enter` to insert a newline.
//
// VS Code::
//
// Add the following to `keybindings.json`:
// [source,json]
// ----
// {
//   "key": "Enter",
//   "command": "rust-analyzer.onEnter",
//   "when": "editorTextFocus && !suggestWidgetVisible && editorLangId == rust"
// }
// ----
//
// When using the Vim plugin:
// [source,json]
// ----
// {
//   "key": "Enter",
//   "command": "rust-analyzer.onEnter",
//   "when": "editorTextFocus && !suggestWidgetVisible && editorLangId == rust && vim.mode == 'Insert'"
// }
// ----
//
// image::https://user-images.githubusercontent.com/48062697/113065578-04c21800-91b1-11eb-82b8-22b8c481e645.gif[]
pub(crate) fn on_enter(db: &RootDatabase, position: FilePosition) -> Option<TextEdit> {
    let parse = db.parse(position.file_id);
    let file = parse.tree();
    let token = file.syntax().token_at_offset(position.offset).left_biased()?;

    if let Some(comment) = ast::Comment::cast(token.clone()) {
        return on_enter_in_comment(&comment, &file, position.offset);
    }

    if token.kind() == L_CURLY {
        // Typing enter after the `{` of a block expression, where the `}` is on the same line
        if let Some(edit) = find_node_at_offset(file.syntax(), position.offset - TextSize::of('{'))
            .and_then(|block| on_enter_in_block(block, position))
        {
            cov_mark::hit!(indent_block_contents);
            return Some(edit);
        }

        // Typing enter after the `{` of a use tree list.
        if let Some(edit) = find_node_at_offset(file.syntax(), position.offset - TextSize::of('{'))
            .and_then(|list| on_enter_in_use_tree_list(list, position))
        {
            cov_mark::hit!(indent_block_contents);
            return Some(edit);
        }
    }

    None
}

fn on_enter_in_comment(
    comment: &ast::Comment,
    file: &ast::SourceFile,
    offset: TextSize,
) -> Option<TextEdit> {
    if comment.kind().shape.is_block() {
        return None;
    }

    let prefix = comment.prefix();
    let comment_range = comment.syntax().text_range();
    if offset < comment_range.start() + TextSize::of(prefix) {
        return None;
    }

    let mut remove_trailing_whitespace = false;
    // Continuing single-line non-doc comments (like this one :) ) is annoying
    if prefix == "//" && comment_range.end() == offset {
        if comment.text().ends_with(' ') {
            cov_mark::hit!(continues_end_of_line_comment_with_space);
            remove_trailing_whitespace = true;
        } else if !followed_by_comment(comment) {
            return None;
        }
    }

    let indent = node_indent(file, comment.syntax())?;
    let inserted = format!("\n{}{} $0", indent, prefix);
    let delete = if remove_trailing_whitespace {
        let trimmed_len = comment.text().trim_end().len() as u32;
        let trailing_whitespace_len = comment.text().len() as u32 - trimmed_len;
        TextRange::new(offset - TextSize::from(trailing_whitespace_len), offset)
    } else {
        TextRange::empty(offset)
    };
    let edit = TextEdit::replace(delete, inserted);
    Some(edit)
}

fn on_enter_in_block(block: ast::BlockExpr, position: FilePosition) -> Option<TextEdit> {
    let contents = block_contents(&block)?;

    if block.syntax().text().contains_char('\n') {
        return None;
    }

    let indent = IndentLevel::from_node(block.syntax());
    let mut edit = TextEdit::insert(position.offset, format!("\n{}$0", indent + 1));
    edit.union(TextEdit::insert(contents.text_range().end(), format!("\n{}", indent))).ok()?;
    Some(edit)
}

fn on_enter_in_use_tree_list(list: ast::UseTreeList, position: FilePosition) -> Option<TextEdit> {
    if list.syntax().text().contains_char('\n') {
        return None;
    }

    let indent = IndentLevel::from_node(list.syntax());
    let mut edit = TextEdit::insert(position.offset, format!("\n{}$0", indent + 1));
    edit.union(TextEdit::insert(
        list.r_curly_token()?.text_range().start(),
        format!("\n{}", indent),
    ))
    .ok()?;
    Some(edit)
}

fn block_contents(block: &ast::BlockExpr) -> Option<SyntaxNode> {
    let mut node = block.tail_expr().map(|e| e.syntax().clone());

    for stmt in block.statements() {
        if node.is_some() {
            // More than 1 node in the block
            return None;
        }

        node = Some(stmt.syntax().clone());
    }

    node
}

fn followed_by_comment(comment: &ast::Comment) -> bool {
    let ws = match comment.syntax().next_token().and_then(ast::Whitespace::cast) {
        Some(it) => it,
        None => return false,
    };
    if ws.spans_multiple_lines() {
        return false;
    }
    ws.syntax().next_token().and_then(ast::Comment::cast).is_some()
}

fn node_indent(file: &SourceFile, token: &SyntaxToken) -> Option<SmolStr> {
    let ws = match file.syntax().token_at_offset(token.text_range().start()) {
        TokenAtOffset::Between(l, r) => {
            assert!(r == *token);
            l
        }
        TokenAtOffset::Single(n) => {
            assert!(n == *token);
            return Some("".into());
        }
        TokenAtOffset::None => unreachable!(),
    };
    if ws.kind() != WHITESPACE {
        return None;
    }
    let text = ws.text();
    let pos = text.rfind('\n').map(|it| it + 1).unwrap_or(0);
    Some(text[pos..].into())
}

#[cfg(test)]
mod tests {
    use stdx::trim_indent;
    use test_utils::assert_eq_text;

    use crate::fixture;

    fn apply_on_enter(before: &str) -> Option<String> {
        let (analysis, position) = fixture::position(before);
        let result = analysis.on_enter(position).unwrap()?;

        let mut actual = analysis.file_text(position.file_id).unwrap().to_string();
        result.apply(&mut actual);
        Some(actual)
    }

    fn do_check(ra_fixture_before: &str, ra_fixture_after: &str) {
        let ra_fixture_after = &trim_indent(ra_fixture_after);
        let actual = apply_on_enter(ra_fixture_before).unwrap();
        assert_eq_text!(ra_fixture_after, &actual);
    }

    fn do_check_noop(ra_fixture_text: &str) {
        assert!(apply_on_enter(ra_fixture_text).is_none())
    }

    #[test]
    fn continues_doc_comment() {
        do_check(
            r"
/// Some docs$0
fn foo() {
}
",
            r"
/// Some docs
/// $0
fn foo() {
}
",
        );

        do_check(
            r"
impl S {
    /// Some$0 docs.
    fn foo() {}
}
",
            r"
impl S {
    /// Some
    /// $0 docs.
    fn foo() {}
}
",
        );

        do_check(
            r"
///$0 Some docs
fn foo() {
}
",
            r"
///
/// $0 Some docs
fn foo() {
}
",
        );
    }

    #[test]
    fn does_not_continue_before_doc_comment() {
        do_check_noop(r"$0//! docz");
    }

    #[test]
    fn continues_another_doc_comment() {
        do_check(
            r#"
fn main() {
    //! Documentation for$0 on enter
    let x = 1 + 1;
}
"#,
            r#"
fn main() {
    //! Documentation for
    //! $0 on enter
    let x = 1 + 1;
}
"#,
        );
    }

    #[test]
    fn continues_code_comment_in_the_middle_of_line() {
        do_check(
            r"
fn main() {
    // Fix$0 me
    let x = 1 + 1;
}
",
            r"
fn main() {
    // Fix
    // $0 me
    let x = 1 + 1;
}
",
        );
    }

    #[test]
    fn continues_code_comment_in_the_middle_several_lines() {
        do_check(
            r"
fn main() {
    // Fix$0
    // me
    let x = 1 + 1;
}
",
            r"
fn main() {
    // Fix
    // $0
    // me
    let x = 1 + 1;
}
",
        );
    }

    #[test]
    fn does_not_continue_end_of_line_comment() {
        do_check_noop(
            r"
fn main() {
    // Fix me$0
    let x = 1 + 1;
}
",
        );
    }

    #[test]
    fn continues_end_of_line_comment_with_space() {
        cov_mark::check!(continues_end_of_line_comment_with_space);
        do_check(
            r#"
fn main() {
    // Fix me $0
    let x = 1 + 1;
}
"#,
            r#"
fn main() {
    // Fix me
    // $0
    let x = 1 + 1;
}
"#,
        );
    }

    #[test]
    fn trims_all_trailing_whitespace() {
        do_check(
            "
fn main() {
    // Fix me  \t\t   $0
    let x = 1 + 1;
}
",
            "
fn main() {
    // Fix me
    // $0
    let x = 1 + 1;
}
",
        );
    }

    #[test]
    fn indents_fn_body_block() {
        cov_mark::check!(indent_block_contents);
        do_check(
            r#"
fn f() {$0()}
        "#,
            r#"
fn f() {
    $0()
}
        "#,
        );
    }

    #[test]
    fn indents_block_expr() {
        do_check(
            r#"
fn f() {
    let x = {$0()};
}
        "#,
            r#"
fn f() {
    let x = {
        $0()
    };
}
        "#,
        );
    }

    #[test]
    fn indents_match_arm() {
        do_check(
            r#"
fn f() {
    match 6 {
        1 => {$0f()},
        _ => (),
    }
}
        "#,
            r#"
fn f() {
    match 6 {
        1 => {
            $0f()
        },
        _ => (),
    }
}
        "#,
        );
    }

    #[test]
    fn indents_block_with_statement() {
        do_check(
            r#"
fn f() {$0a = b}
        "#,
            r#"
fn f() {
    $0a = b
}
        "#,
        );
        do_check(
            r#"
fn f() {$0fn f() {}}
        "#,
            r#"
fn f() {
    $0fn f() {}
}
        "#,
        );
    }

    #[test]
    fn indents_nested_blocks() {
        do_check(
            r#"
fn f() {$0{}}
        "#,
            r#"
fn f() {
    $0{}
}
        "#,
        );
    }

    #[test]
    fn does_not_indent_empty_block() {
        do_check_noop(
            r#"
fn f() {$0}
        "#,
        );
        do_check_noop(
            r#"
fn f() {{$0}}
        "#,
        );
    }

    #[test]
    fn does_not_indent_block_with_too_much_content() {
        do_check_noop(
            r#"
fn f() {$0 a = b; ()}
        "#,
        );
        do_check_noop(
            r#"
fn f() {$0 a = b; a = b; }
        "#,
        );
    }

    #[test]
    fn does_not_indent_multiline_block() {
        do_check_noop(
            r#"
fn f() {$0
}
        "#,
        );
        do_check_noop(
            r#"
fn f() {$0

}
        "#,
        );
    }

    #[test]
    fn indents_use_tree_list() {
        do_check(
            r#"
use crate::{$0};
            "#,
            r#"
use crate::{
    $0
};
            "#,
        );
        do_check(
            r#"
use crate::{$0Object, path::to::OtherThing};
            "#,
            r#"
use crate::{
    $0Object, path::to::OtherThing
};
            "#,
        );
        do_check(
            r#"
use {crate::{$0Object, path::to::OtherThing}};
            "#,
            r#"
use {crate::{
    $0Object, path::to::OtherThing
}};
            "#,
        );
        do_check(
            r#"
use {
    crate::{$0Object, path::to::OtherThing}
};
            "#,
            r#"
use {
    crate::{
        $0Object, path::to::OtherThing
    }
};
            "#,
        );
    }

    #[test]
    fn does_not_indent_use_tree_list_when_not_at_curly_brace() {
        do_check_noop(
            r#"
use path::{Thing$0};
            "#,
        );
    }

    #[test]
    fn does_not_indent_use_tree_list_without_curly_braces() {
        do_check_noop(
            r#"
use path::Thing$0;
            "#,
        );
        do_check_noop(
            r#"
use path::$0Thing;
            "#,
        );
        do_check_noop(
            r#"
use path::Thing$0};
            "#,
        );
        do_check_noop(
            r#"
use path::{$0Thing;
            "#,
        );
    }

    #[test]
    fn does_not_indent_multiline_use_tree_list() {
        do_check_noop(
            r#"
use path::{$0
    Thing
};
            "#,
        );
    }
}
<|endoftext|>


<commit_before>//! Handles the `Enter` key press. At the momently, this only continues
//! comments, but should handle indent some time in the future as well.

use std::sync::Arc;

use ide_db::RootDatabase;
use ide_db::{
    base_db::{FilePosition, SourceDatabase},
    line_index::LineIndex,
    LineIndexDatabase,
};
use syntax::{
    algo::find_node_at_offset,
    ast::{self, edit::IndentLevel, AstToken},
    AstNode, SmolStr, SourceFile,
    SyntaxKind::*,
    SyntaxNode, SyntaxToken, TextRange, TextSize, TokenAtOffset,
};

use text_edit::TextEdit;

// Feature: On Enter
//
// rust-analyzer can override kbd:[Enter] key to make it smarter:
//
// - kbd:[Enter] inside triple-slash comments automatically inserts `///`
// - kbd:[Enter] in the middle or after a trailing space in `//` inserts `//`
// - kbd:[Enter] inside `//!` doc comments automatically inserts `//!`
// - kbd:[Enter] after `{` indents contents and closing `}` of single-line block
//
// This action needs to be assigned to shortcut explicitly.
//
// VS Code::
//
// Add the following to `keybindings.json`:
// [source,json]
// ----
// {
//   "key": "Enter",
//   "command": "rust-analyzer.onEnter",
//   "when": "editorTextFocus && !suggestWidgetVisible && editorLangId == rust"
// }
// ----
//
// image::https://user-images.githubusercontent.com/48062697/113065578-04c21800-91b1-11eb-82b8-22b8c481e645.gif[]
pub(crate) fn on_enter(db: &RootDatabase, position: FilePosition) -> Option<TextEdit> {
    let parse = db.parse(position.file_id);
    let file = parse.tree();
    let token = file.syntax().token_at_offset(position.offset).left_biased()?;

    if let Some(comment) = ast::Comment::cast(token.clone()) {
        return on_enter_in_comment(&comment, &file, position.offset);
    }

    if token.kind() == L_CURLY {
        // Typing enter after the `{` of a block expression, where the `}` is on the same line
        if let Some(edit) = find_node_at_offset(file.syntax(), position.offset - TextSize::of('{'))
            .and_then(|block| on_enter_in_block(db, block, position))
        {
            return Some(edit);
        }
    }

    None
}

fn on_enter_in_comment(
    comment: &ast::Comment,
    file: &ast::SourceFile,
    offset: TextSize,
) -> Option<TextEdit> {
    if comment.kind().shape.is_block() {
        return None;
    }

    let prefix = comment.prefix();
    let comment_range = comment.syntax().text_range();
    if offset < comment_range.start() + TextSize::of(prefix) {
        return None;
    }

    let mut remove_trailing_whitespace = false;
    // Continuing single-line non-doc comments (like this one :) ) is annoying
    if prefix == "//" && comment_range.end() == offset {
        if comment.text().ends_with(' ') {
            cov_mark::hit!(continues_end_of_line_comment_with_space);
            remove_trailing_whitespace = true;
        } else if !followed_by_comment(&comment) {
            return None;
        }
    }

    let indent = node_indent(&file, comment.syntax())?;
    let inserted = format!("\n{}{} $0", indent, prefix);
    let delete = if remove_trailing_whitespace {
        let trimmed_len = comment.text().trim_end().len() as u32;
        let trailing_whitespace_len = comment.text().len() as u32 - trimmed_len;
        TextRange::new(offset - TextSize::from(trailing_whitespace_len), offset)
    } else {
        TextRange::empty(offset)
    };
    let edit = TextEdit::replace(delete, inserted);
    Some(edit)
}

fn on_enter_in_block(
    db: &RootDatabase,
    block: ast::BlockExpr,
    position: FilePosition,
) -> Option<TextEdit> {
    let contents = block_contents(&block)?;

    let line_index: Arc<LineIndex> = db.line_index(position.file_id);
    let (open, close) = (block.l_curly_token()?, block.r_curly_token()?);
    let start = line_index.line_col(open.text_range().start()).line;
    let end = line_index.line_col(close.text_range().end()).line;
    if start != end {
        return None;
    }

    let indent = IndentLevel::from_node(block.syntax());
    let mut edit = TextEdit::insert(position.offset, format!("\n{}$0", indent + 1));
    edit.union(TextEdit::insert(contents.text_range().end(), format!("\n{}", indent))).ok()?;
    Some(edit)
}

fn block_contents(block: &ast::BlockExpr) -> Option<SyntaxNode> {
    let mut node = block.tail_expr().map(|e| e.syntax().clone());

    for stmt in block.statements() {
        if node.is_some() {
            // More than 1 node in the block
            return None;
        }

        node = Some(stmt.syntax().clone());
    }

    node
}

fn followed_by_comment(comment: &ast::Comment) -> bool {
    let ws = match comment.syntax().next_token().and_then(ast::Whitespace::cast) {
        Some(it) => it,
        None => return false,
    };
    if ws.spans_multiple_lines() {
        return false;
    }
    ws.syntax().next_token().and_then(ast::Comment::cast).is_some()
}

fn node_indent(file: &SourceFile, token: &SyntaxToken) -> Option<SmolStr> {
    let ws = match file.syntax().token_at_offset(token.text_range().start()) {
        TokenAtOffset::Between(l, r) => {
            assert!(r == *token);
            l
        }
        TokenAtOffset::Single(n) => {
            assert!(n == *token);
            return Some("".into());
        }
        TokenAtOffset::None => unreachable!(),
    };
    if ws.kind() != WHITESPACE {
        return None;
    }
    let text = ws.text();
    let pos = text.rfind('\n').map(|it| it + 1).unwrap_or(0);
    Some(text[pos..].into())
}

#[cfg(test)]
mod tests {
    use stdx::trim_indent;
    use test_utils::assert_eq_text;

    use crate::fixture;

    fn apply_on_enter(before: &str) -> Option<String> {
        let (analysis, position) = fixture::position(&before);
        let result = analysis.on_enter(position).unwrap()?;

        let mut actual = analysis.file_text(position.file_id).unwrap().to_string();
        result.apply(&mut actual);
        Some(actual)
    }

    fn do_check(ra_fixture_before: &str, ra_fixture_after: &str) {
        let ra_fixture_after = &trim_indent(ra_fixture_after);
        let actual = apply_on_enter(ra_fixture_before).unwrap();
        assert_eq_text!(ra_fixture_after, &actual);
    }

    fn do_check_noop(ra_fixture_text: &str) {
        assert!(apply_on_enter(ra_fixture_text).is_none())
    }

    #[test]
    fn continues_doc_comment() {
        do_check(
            r"
/// Some docs$0
fn foo() {
}
",
            r"
/// Some docs
/// $0
fn foo() {
}
",
        );

        do_check(
            r"
impl S {
    /// Some$0 docs.
    fn foo() {}
}
",
            r"
impl S {
    /// Some
    /// $0 docs.
    fn foo() {}
}
",
        );

        do_check(
            r"
///$0 Some docs
fn foo() {
}
",
            r"
///
/// $0 Some docs
fn foo() {
}
",
        );
    }

    #[test]
    fn does_not_continue_before_doc_comment() {
        do_check_noop(r"$0//! docz");
    }

    #[test]
    fn continues_another_doc_comment() {
        do_check(
            r#"
fn main() {
    //! Documentation for$0 on enter
    let x = 1 + 1;
}
"#,
            r#"
fn main() {
    //! Documentation for
    //! $0 on enter
    let x = 1 + 1;
}
"#,
        );
    }

    #[test]
    fn continues_code_comment_in_the_middle_of_line() {
        do_check(
            r"
fn main() {
    // Fix$0 me
    let x = 1 + 1;
}
",
            r"
fn main() {
    // Fix
    // $0 me
    let x = 1 + 1;
}
",
        );
    }

    #[test]
    fn continues_code_comment_in_the_middle_several_lines() {
        do_check(
            r"
fn main() {
    // Fix$0
    // me
    let x = 1 + 1;
}
",
            r"
fn main() {
    // Fix
    // $0
    // me
    let x = 1 + 1;
}
",
        );
    }

    #[test]
    fn does_not_continue_end_of_line_comment() {
        do_check_noop(
            r"
fn main() {
    // Fix me$0
    let x = 1 + 1;
}
",
        );
    }

    #[test]
    fn continues_end_of_line_comment_with_space() {
        cov_mark::check!(continues_end_of_line_comment_with_space);
        do_check(
            r#"
fn main() {
    // Fix me $0
    let x = 1 + 1;
}
"#,
            r#"
fn main() {
    // Fix me
    // $0
    let x = 1 + 1;
}
"#,
        );
    }

    #[test]
    fn trims_all_trailing_whitespace() {
        do_check(
            "
fn main() {
    // Fix me  \t\t   $0
    let x = 1 + 1;
}
",
            "
fn main() {
    // Fix me
    // $0
    let x = 1 + 1;
}
",
        );
    }

    #[test]
    fn indents_fn_body_block() {
        do_check(
            r#"
fn f() {$0()}
        "#,
            r#"
fn f() {
    $0()
}
        "#,
        );
    }

    #[test]
    fn indents_block_expr() {
        do_check(
            r#"
fn f() {
    let x = {$0()};
}
        "#,
            r#"
fn f() {
    let x = {
        $0()
    };
}
        "#,
        );
    }

    #[test]
    fn indents_match_arm() {
        do_check(
            r#"
fn f() {
    match 6 {
        1 => {$0f()},
        _ => (),
    }
}
        "#,
            r#"
fn f() {
    match 6 {
        1 => {
            $0f()
        },
        _ => (),
    }
}
        "#,
        );
    }

    #[test]
    fn indents_block_with_statement() {
        do_check(
            r#"
fn f() {$0a = b}
        "#,
            r#"
fn f() {
    $0a = b
}
        "#,
        );
        do_check(
            r#"
fn f() {$0fn f() {}}
        "#,
            r#"
fn f() {
    $0fn f() {}
}
        "#,
        );
    }

    #[test]
    fn indents_nested_blocks() {
        do_check(
            r#"
fn f() {$0{}}
        "#,
            r#"
fn f() {
    $0{}
}
        "#,
        );
    }

    #[test]
    fn does_not_indent_empty_block() {
        do_check_noop(
            r#"
fn f() {$0}
        "#,
        );
        do_check_noop(
            r#"
fn f() {{$0}}
        "#,
        );
    }

    #[test]
    fn does_not_indent_block_with_too_much_content() {
        do_check_noop(
            r#"
fn f() {$0 a = b; ()}
        "#,
        );
        do_check_noop(
            r#"
fn f() {$0 a = b; a = b; }
        "#,
        );
    }

    #[test]
    fn does_not_indent_multiline_block() {
        do_check_noop(
            r#"
fn f() {$0
}
        "#,
        );
        do_check_noop(
            r#"
fn f() {$0

}
        "#,
        );
    }
}
<commit_msg>Simplify multiline check<commit_after>//! Handles the `Enter` key press. At the momently, this only continues
//! comments, but should handle indent some time in the future as well.

use ide_db::base_db::{FilePosition, SourceDatabase};
use ide_db::RootDatabase;
use syntax::{
    algo::find_node_at_offset,
    ast::{self, edit::IndentLevel, AstToken},
    AstNode, SmolStr, SourceFile,
    SyntaxKind::*,
    SyntaxNode, SyntaxToken, TextRange, TextSize, TokenAtOffset,
};

use text_edit::TextEdit;

// Feature: On Enter
//
// rust-analyzer can override kbd:[Enter] key to make it smarter:
//
// - kbd:[Enter] inside triple-slash comments automatically inserts `///`
// - kbd:[Enter] in the middle or after a trailing space in `//` inserts `//`
// - kbd:[Enter] inside `//!` doc comments automatically inserts `//!`
// - kbd:[Enter] after `{` indents contents and closing `}` of single-line block
//
// This action needs to be assigned to shortcut explicitly.
//
// VS Code::
//
// Add the following to `keybindings.json`:
// [source,json]
// ----
// {
//   "key": "Enter",
//   "command": "rust-analyzer.onEnter",
//   "when": "editorTextFocus && !suggestWidgetVisible && editorLangId == rust"
// }
// ----
//
// image::https://user-images.githubusercontent.com/48062697/113065578-04c21800-91b1-11eb-82b8-22b8c481e645.gif[]
pub(crate) fn on_enter(db: &RootDatabase, position: FilePosition) -> Option<TextEdit> {
    let parse = db.parse(position.file_id);
    let file = parse.tree();
    let token = file.syntax().token_at_offset(position.offset).left_biased()?;

    if let Some(comment) = ast::Comment::cast(token.clone()) {
        return on_enter_in_comment(&comment, &file, position.offset);
    }

    if token.kind() == L_CURLY {
        // Typing enter after the `{` of a block expression, where the `}` is on the same line
        if let Some(edit) = find_node_at_offset(file.syntax(), position.offset - TextSize::of('{'))
            .and_then(|block| on_enter_in_block(block, position))
        {
            return Some(edit);
        }
    }

    None
}

fn on_enter_in_comment(
    comment: &ast::Comment,
    file: &ast::SourceFile,
    offset: TextSize,
) -> Option<TextEdit> {
    if comment.kind().shape.is_block() {
        return None;
    }

    let prefix = comment.prefix();
    let comment_range = comment.syntax().text_range();
    if offset < comment_range.start() + TextSize::of(prefix) {
        return None;
    }

    let mut remove_trailing_whitespace = false;
    // Continuing single-line non-doc comments (like this one :) ) is annoying
    if prefix == "//" && comment_range.end() == offset {
        if comment.text().ends_with(' ') {
            cov_mark::hit!(continues_end_of_line_comment_with_space);
            remove_trailing_whitespace = true;
        } else if !followed_by_comment(&comment) {
            return None;
        }
    }

    let indent = node_indent(&file, comment.syntax())?;
    let inserted = format!("\n{}{} $0", indent, prefix);
    let delete = if remove_trailing_whitespace {
        let trimmed_len = comment.text().trim_end().len() as u32;
        let trailing_whitespace_len = comment.text().len() as u32 - trimmed_len;
        TextRange::new(offset - TextSize::from(trailing_whitespace_len), offset)
    } else {
        TextRange::empty(offset)
    };
    let edit = TextEdit::replace(delete, inserted);
    Some(edit)
}

fn on_enter_in_block(block: ast::BlockExpr, position: FilePosition) -> Option<TextEdit> {
    let contents = block_contents(&block)?;

    if block.syntax().text().contains_char('\n') {
        return None;
    }

    let indent = IndentLevel::from_node(block.syntax());
    let mut edit = TextEdit::insert(position.offset, format!("\n{}$0", indent + 1));
    edit.union(TextEdit::insert(contents.text_range().end(), format!("\n{}", indent))).ok()?;
    Some(edit)
}

fn block_contents(block: &ast::BlockExpr) -> Option<SyntaxNode> {
    let mut node = block.tail_expr().map(|e| e.syntax().clone());

    for stmt in block.statements() {
        if node.is_some() {
            // More than 1 node in the block
            return None;
        }

        node = Some(stmt.syntax().clone());
    }

    node
}

fn followed_by_comment(comment: &ast::Comment) -> bool {
    let ws = match comment.syntax().next_token().and_then(ast::Whitespace::cast) {
        Some(it) => it,
        None => return false,
    };
    if ws.spans_multiple_lines() {
        return false;
    }
    ws.syntax().next_token().and_then(ast::Comment::cast).is_some()
}

fn node_indent(file: &SourceFile, token: &SyntaxToken) -> Option<SmolStr> {
    let ws = match file.syntax().token_at_offset(token.text_range().start()) {
        TokenAtOffset::Between(l, r) => {
            assert!(r == *token);
            l
        }
        TokenAtOffset::Single(n) => {
            assert!(n == *token);
            return Some("".into());
        }
        TokenAtOffset::None => unreachable!(),
    };
    if ws.kind() != WHITESPACE {
        return None;
    }
    let text = ws.text();
    let pos = text.rfind('\n').map(|it| it + 1).unwrap_or(0);
    Some(text[pos..].into())
}

#[cfg(test)]
mod tests {
    use stdx::trim_indent;
    use test_utils::assert_eq_text;

    use crate::fixture;

    fn apply_on_enter(before: &str) -> Option<String> {
        let (analysis, position) = fixture::position(&before);
        let result = analysis.on_enter(position).unwrap()?;

        let mut actual = analysis.file_text(position.file_id).unwrap().to_string();
        result.apply(&mut actual);
        Some(actual)
    }

    fn do_check(ra_fixture_before: &str, ra_fixture_after: &str) {
        let ra_fixture_after = &trim_indent(ra_fixture_after);
        let actual = apply_on_enter(ra_fixture_before).unwrap();
        assert_eq_text!(ra_fixture_after, &actual);
    }

    fn do_check_noop(ra_fixture_text: &str) {
        assert!(apply_on_enter(ra_fixture_text).is_none())
    }

    #[test]
    fn continues_doc_comment() {
        do_check(
            r"
/// Some docs$0
fn foo() {
}
",
            r"
/// Some docs
/// $0
fn foo() {
}
",
        );

        do_check(
            r"
impl S {
    /// Some$0 docs.
    fn foo() {}
}
",
            r"
impl S {
    /// Some
    /// $0 docs.
    fn foo() {}
}
",
        );

        do_check(
            r"
///$0 Some docs
fn foo() {
}
",
            r"
///
/// $0 Some docs
fn foo() {
}
",
        );
    }

    #[test]
    fn does_not_continue_before_doc_comment() {
        do_check_noop(r"$0//! docz");
    }

    #[test]
    fn continues_another_doc_comment() {
        do_check(
            r#"
fn main() {
    //! Documentation for$0 on enter
    let x = 1 + 1;
}
"#,
            r#"
fn main() {
    //! Documentation for
    //! $0 on enter
    let x = 1 + 1;
}
"#,
        );
    }

    #[test]
    fn continues_code_comment_in_the_middle_of_line() {
        do_check(
            r"
fn main() {
    // Fix$0 me
    let x = 1 + 1;
}
",
            r"
fn main() {
    // Fix
    // $0 me
    let x = 1 + 1;
}
",
        );
    }

    #[test]
    fn continues_code_comment_in_the_middle_several_lines() {
        do_check(
            r"
fn main() {
    // Fix$0
    // me
    let x = 1 + 1;
}
",
            r"
fn main() {
    // Fix
    // $0
    // me
    let x = 1 + 1;
}
",
        );
    }

    #[test]
    fn does_not_continue_end_of_line_comment() {
        do_check_noop(
            r"
fn main() {
    // Fix me$0
    let x = 1 + 1;
}
",
        );
    }

    #[test]
    fn continues_end_of_line_comment_with_space() {
        cov_mark::check!(continues_end_of_line_comment_with_space);
        do_check(
            r#"
fn main() {
    // Fix me $0
    let x = 1 + 1;
}
"#,
            r#"
fn main() {
    // Fix me
    // $0
    let x = 1 + 1;
}
"#,
        );
    }

    #[test]
    fn trims_all_trailing_whitespace() {
        do_check(
            "
fn main() {
    // Fix me  \t\t   $0
    let x = 1 + 1;
}
",
            "
fn main() {
    // Fix me
    // $0
    let x = 1 + 1;
}
",
        );
    }

    #[test]
    fn indents_fn_body_block() {
        do_check(
            r#"
fn f() {$0()}
        "#,
            r#"
fn f() {
    $0()
}
        "#,
        );
    }

    #[test]
    fn indents_block_expr() {
        do_check(
            r#"
fn f() {
    let x = {$0()};
}
        "#,
            r#"
fn f() {
    let x = {
        $0()
    };
}
        "#,
        );
    }

    #[test]
    fn indents_match_arm() {
        do_check(
            r#"
fn f() {
    match 6 {
        1 => {$0f()},
        _ => (),
    }
}
        "#,
            r#"
fn f() {
    match 6 {
        1 => {
            $0f()
        },
        _ => (),
    }
}
        "#,
        );
    }

    #[test]
    fn indents_block_with_statement() {
        do_check(
            r#"
fn f() {$0a = b}
        "#,
            r#"
fn f() {
    $0a = b
}
        "#,
        );
        do_check(
            r#"
fn f() {$0fn f() {}}
        "#,
            r#"
fn f() {
    $0fn f() {}
}
        "#,
        );
    }

    #[test]
    fn indents_nested_blocks() {
        do_check(
            r#"
fn f() {$0{}}
        "#,
            r#"
fn f() {
    $0{}
}
        "#,
        );
    }

    #[test]
    fn does_not_indent_empty_block() {
        do_check_noop(
            r#"
fn f() {$0}
        "#,
        );
        do_check_noop(
            r#"
fn f() {{$0}}
        "#,
        );
    }

    #[test]
    fn does_not_indent_block_with_too_much_content() {
        do_check_noop(
            r#"
fn f() {$0 a = b; ()}
        "#,
        );
        do_check_noop(
            r#"
fn f() {$0 a = b; a = b; }
        "#,
        );
    }

    #[test]
    fn does_not_indent_multiline_block() {
        do_check_noop(
            r#"
fn f() {$0
}
        "#,
        );
        do_check_noop(
            r#"
fn f() {$0

}
        "#,
        );
    }
}
<|endoftext|>