<commit_before>use types::*;

use nom::{be_i8, be_u8, be_i16, be_u16, be_i32, be_u32, be_i64, be_u64, be_f32, be_f64};

/* FIXME: we convert to a Some so that nom can happily handle some _ case which would otherwise fail as we're exhaustive */
named!(pub parse_value<AMQPValue>,                 switch!(map!(parse_type, |t| Some(t)),
    Some(AMQPType::Boolean)        => map!(call!(parse_boolean),          |b| AMQPValue::Boolean(b))        |
    Some(AMQPType::ShortShortInt)  => map!(call!(parse_short_short_int),  |i| AMQPValue::ShortShortInt(i))  |
    Some(AMQPType::ShortShortUInt) => map!(call!(parse_short_short_uint), |u| AMQPValue::ShortShortUInt(u)) |
    Some(AMQPType::ShortInt)       => map!(call!(parse_short_int),        |i| AMQPValue::ShortInt(i))       |
    Some(AMQPType::ShortUInt)      => map!(call!(parse_short_uint),       |u| AMQPValue::ShortUInt(u))      |
    Some(AMQPType::LongInt)        => map!(call!(parse_long_int),         |i| AMQPValue::LongInt(i))        |
    Some(AMQPType::LongUInt)       => map!(call!(parse_long_uint),        |u| AMQPValue::LongUInt(u))       |
    Some(AMQPType::LongLongInt)    => map!(call!(parse_long_long_int),    |i| AMQPValue::LongLongInt(i))    |
    Some(AMQPType::LongLongUInt)   => map!(call!(parse_long_long_uint),   |u| AMQPValue::LongLongUInt(u))   |
    Some(AMQPType::Float)          => map!(call!(parse_float),            |f| AMQPValue::Float(f))          |
    Some(AMQPType::Double)         => map!(call!(parse_double),           |d| AMQPValue::Double(d))         |
    Some(AMQPType::DecimalValue)   => map!(call!(parse_decimal_value),    |d| AMQPValue::DecimalValue(d))   |
    Some(AMQPType::ShortString)    => map!(call!(parse_short_string),     |s| AMQPValue::ShortString(s))    |
    Some(AMQPType::LongString)     => map!(call!(parse_long_string),      |s| AMQPValue::LongString(s))     |
    Some(AMQPType::FieldArray)     => map!(call!(parse_field_array),      |a| AMQPValue::FieldArray(a))     |
    Some(AMQPType::Timestamp)      => map!(call!(parse_timestamp),        |t| AMQPValue::Timestamp(t))      |
    Some(AMQPType::FieldTable)     => map!(call!(parse_field_table),      |t| AMQPValue::FieldTable(t))     |
    Some(AMQPType::Void)           => value!(AMQPValue::Void)
));

named!(pub parse_type<AMQPType>,                   map_opt!(be_u8, |t| AMQPType::from_id(t as char)));

named!(pub parse_boolean<Boolean>,                 map!(be_u8, |b| b != 0));
named!(pub parse_short_short_int<ShortShortInt>,   call!(be_i8));
named!(pub parse_short_short_uint<ShortShortUInt>, call!(be_u8));
named!(pub parse_short_int<ShortInt>,              call!(be_i16));
named!(pub parse_short_uint<ShortUInt>,            call!(be_u16));
named!(pub parse_long_int<LongInt>,                call!(be_i32));
named!(pub parse_long_uint<LongUInt>,              call!(be_u32));
named!(pub parse_long_long_int<LongLongInt>,       call!(be_i64));
named!(pub parse_long_long_uint<LongLongUInt>,     call!(be_u64));
named!(pub parse_float<Float>,                     call!(be_f32));
named!(pub parse_double<Double>,                   call!(be_f64));
named!(pub parse_decimal_value<DecimalValue>,      do_parse!(scale: parse_short_short_uint >> value: parse_long_uint >> (DecimalValue { scale: scale, value: value, })));
named!(pub parse_short_string<ShortString>,        do_parse!(length: parse_short_short_uint >> s: take_str!(length) >> (s.to_string())));
named!(pub parse_long_string<LongString>,          do_parse!(length: parse_long_uint >> s: take_str!(length) >> (s.to_string())));
named!(pub parse_field_array<FieldArray>,          do_parse!(length: parse_long_int >> array: count!(parse_value, length as usize) >> (array)));
named!(pub parse_timestamp<Timestamp>,             call!(parse_long_long_uint));
named!(pub parse_field_table<FieldTable>,          do_parse!(length: parse_long_uint >> table: flat_map!(take!(length as usize), fold_many0!(complete!(pair!(parse_short_string, parse_value)), FieldTable::new(), |mut acc: FieldTable, (key, value)| {
    acc.insert(key, value);
    acc
})) >> (table)));

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

    use nom::IResult;

    const EMPTY: &'static [u8] = b"";

    #[test]
    fn test_parse_boolean() {
        assert_eq!(parse_boolean(&[0]), IResult::Done(EMPTY, false));
        assert_eq!(parse_boolean(&[1]), IResult::Done(EMPTY, true));
    }

    #[test]
    fn test_parse_short_short_uint() {
        assert_eq!(parse_short_short_uint(&[0]),   IResult::Done(EMPTY, 0));
        assert_eq!(parse_short_short_uint(&[255]), IResult::Done(EMPTY, 255));
    }

    #[test]
    fn test_parse_short_short_int() {
        assert_eq!(parse_short_short_int(&[0]),   IResult::Done(EMPTY, 0));
        assert_eq!(parse_short_short_int(&[255]), IResult::Done(EMPTY, -1));
    }

    #[test]
    fn test_parse_short_uint() {
        assert_eq!(parse_short_uint(&[0,   0]),   IResult::Done(EMPTY, 0));
        assert_eq!(parse_short_uint(&[255, 255]), IResult::Done(EMPTY, 65535));
    }

    #[test]
    fn test_parse_short_int() {
        assert_eq!(parse_short_int(&[0,   0]),   IResult::Done(EMPTY, 0));
        assert_eq!(parse_short_int(&[255, 255]), IResult::Done(EMPTY, -1));
    }

    #[test]
    fn test_parse_long_uint() {
        assert_eq!(parse_long_uint(&[0,   0,   0,   0]),   IResult::Done(EMPTY, 0));
        assert_eq!(parse_long_uint(&[255, 255, 255, 255]), IResult::Done(EMPTY, 4294967295));
    }

    #[test]
    fn test_parse_long_int() {
        assert_eq!(parse_long_int(&[0,   0,   0,   0]),   IResult::Done(EMPTY, 0));
        assert_eq!(parse_long_int(&[255, 255, 255, 255]), IResult::Done(EMPTY, -1));
    }

    #[test]
    fn test_parse_long_long_uint() {
        assert_eq!(parse_long_long_uint(&[0,   0,   0,   0,   0,   0,   0,   0]),   IResult::Done(EMPTY, 0));
        assert_eq!(parse_long_long_uint(&[255, 255, 255, 255, 255, 255, 255, 255]), IResult::Done(EMPTY, 18446744073709551615));
    }

    #[test]
    fn test_parse_long_long_int() {
        assert_eq!(parse_long_long_int(&[0,   0,   0,   0,   0,   0,   0,   0]),   IResult::Done(EMPTY, 0));
        assert_eq!(parse_long_long_int(&[255, 255, 255, 255, 255, 255, 255, 255]), IResult::Done(EMPTY, -1));
    }

    #[test]
    fn test_parse_float() {
        assert_eq!(parse_float(&[0,  0,  0,  0]),  IResult::Done(EMPTY, 0.));
        assert_eq!(parse_float(&[42, 42, 42, 42]), IResult::Done(EMPTY, 0.00000000000015113662));
    }

    #[test]
    fn test_parse_double() {
        assert_eq!(parse_double(&[0,  0,  0,  0,  0,  0,  0,  0]),  IResult::Done(EMPTY, 0.));
        assert_eq!(parse_double(&[42, 42, 42, 42, 42, 42, 42, 42]), IResult::Done(EMPTY, 0.0000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000014260258159703532));
    }

    #[test]
    fn test_parse_decimal_value() {
        assert_eq!(parse_decimal_value(&[0,   0,   0,   0,   0]),   IResult::Done(EMPTY, DecimalValue { scale: 0,   value: 0          }));
        assert_eq!(parse_decimal_value(&[255, 255, 255, 255, 255]), IResult::Done(EMPTY, DecimalValue { scale: 255, value: 4294967295 }));
    }

    #[test]
    fn test_parse_short_string() {
        assert_eq!(parse_short_string(&[0]),                     IResult::Done(EMPTY, ShortString::new()));
        assert_eq!(parse_short_string(&[4, 116, 101, 115, 116]), IResult::Done(EMPTY, "test".to_string()));
    }

    #[test]
    fn test_parse_long_string() {
        assert_eq!(parse_long_string(&[0, 0, 0, 0]),                     IResult::Done(EMPTY, LongString::new()));
        assert_eq!(parse_long_string(&[0, 0, 0, 4, 116, 101, 115, 116]), IResult::Done(EMPTY, "test".to_string()));
    }

    #[test]
    fn test_parse_field_array() {
        assert_eq!(parse_field_array(&[0, 0, 0, 0]),                                 IResult::Done(EMPTY, FieldArray::new()));
        assert_eq!(parse_field_array(&[0, 0, 0, 2, 115, 4, 116, 101, 115, 116, 86]), IResult::Done(EMPTY, vec![AMQPValue::ShortString("test".to_string()), AMQPValue::Void]));
    }

    #[test]
    fn test_parse_timestamp() {
        assert_eq!(parse_timestamp(&[0,   0,   0,   0,   0,   0,   0,   0]),   IResult::Done(EMPTY, 0));
        assert_eq!(parse_timestamp(&[255, 255, 255, 255, 255, 255, 255, 255]), IResult::Done(EMPTY, 18446744073709551615));
    }

    #[test]
    fn test_parse_field_table() {
        let mut table = FieldTable::new();
        table.insert("test".to_string(), AMQPValue::ShortString("test".to_string()));
        table.insert("tt".to_string(),   AMQPValue::Void);
        assert_eq!(parse_field_table(&[0, 0, 0, 0]),                                                                      IResult::Done(EMPTY, FieldTable::new()));
        assert_eq!(parse_field_table(&[0, 0, 0, 15, 4, 116, 101, 115, 116, 115, 4, 116, 101, 115, 116, 2, 116, 116, 86]), IResult::Done(EMPTY, table));
    }
}
<commit_msg>types: more parsing tests<commit_after>use types::*;

use nom::{be_i8, be_u8, be_i16, be_u16, be_i32, be_u32, be_i64, be_u64, be_f32, be_f64};

/* FIXME: we convert to a Some so that nom can happily handle some _ case which would otherwise fail as we're exhaustive */
named!(pub parse_value<AMQPValue>,                 switch!(map!(parse_type, |t| Some(t)),
    Some(AMQPType::Boolean)        => map!(call!(parse_boolean),          |b| AMQPValue::Boolean(b))        |
    Some(AMQPType::ShortShortInt)  => map!(call!(parse_short_short_int),  |i| AMQPValue::ShortShortInt(i))  |
    Some(AMQPType::ShortShortUInt) => map!(call!(parse_short_short_uint), |u| AMQPValue::ShortShortUInt(u)) |
    Some(AMQPType::ShortInt)       => map!(call!(parse_short_int),        |i| AMQPValue::ShortInt(i))       |
    Some(AMQPType::ShortUInt)      => map!(call!(parse_short_uint),       |u| AMQPValue::ShortUInt(u))      |
    Some(AMQPType::LongInt)        => map!(call!(parse_long_int),         |i| AMQPValue::LongInt(i))        |
    Some(AMQPType::LongUInt)       => map!(call!(parse_long_uint),        |u| AMQPValue::LongUInt(u))       |
    Some(AMQPType::LongLongInt)    => map!(call!(parse_long_long_int),    |i| AMQPValue::LongLongInt(i))    |
    Some(AMQPType::LongLongUInt)   => map!(call!(parse_long_long_uint),   |u| AMQPValue::LongLongUInt(u))   |
    Some(AMQPType::Float)          => map!(call!(parse_float),            |f| AMQPValue::Float(f))          |
    Some(AMQPType::Double)         => map!(call!(parse_double),           |d| AMQPValue::Double(d))         |
    Some(AMQPType::DecimalValue)   => map!(call!(parse_decimal_value),    |d| AMQPValue::DecimalValue(d))   |
    Some(AMQPType::ShortString)    => map!(call!(parse_short_string),     |s| AMQPValue::ShortString(s))    |
    Some(AMQPType::LongString)     => map!(call!(parse_long_string),      |s| AMQPValue::LongString(s))     |
    Some(AMQPType::FieldArray)     => map!(call!(parse_field_array),      |a| AMQPValue::FieldArray(a))     |
    Some(AMQPType::Timestamp)      => map!(call!(parse_timestamp),        |t| AMQPValue::Timestamp(t))      |
    Some(AMQPType::FieldTable)     => map!(call!(parse_field_table),      |t| AMQPValue::FieldTable(t))     |
    Some(AMQPType::Void)           => value!(AMQPValue::Void)
));

named!(pub parse_type<AMQPType>,                   map_opt!(be_u8, |t| AMQPType::from_id(t as char)));

named!(pub parse_boolean<Boolean>,                 map!(be_u8, |b| b != 0));
named!(pub parse_short_short_int<ShortShortInt>,   call!(be_i8));
named!(pub parse_short_short_uint<ShortShortUInt>, call!(be_u8));
named!(pub parse_short_int<ShortInt>,              call!(be_i16));
named!(pub parse_short_uint<ShortUInt>,            call!(be_u16));
named!(pub parse_long_int<LongInt>,                call!(be_i32));
named!(pub parse_long_uint<LongUInt>,              call!(be_u32));
named!(pub parse_long_long_int<LongLongInt>,       call!(be_i64));
named!(pub parse_long_long_uint<LongLongUInt>,     call!(be_u64));
named!(pub parse_float<Float>,                     call!(be_f32));
named!(pub parse_double<Double>,                   call!(be_f64));
named!(pub parse_decimal_value<DecimalValue>,      do_parse!(scale: parse_short_short_uint >> value: parse_long_uint >> (DecimalValue { scale: scale, value: value, })));
named!(pub parse_short_string<ShortString>,        do_parse!(length: parse_short_short_uint >> s: take_str!(length) >> (s.to_string())));
named!(pub parse_long_string<LongString>,          do_parse!(length: parse_long_uint >> s: take_str!(length) >> (s.to_string())));
named!(pub parse_field_array<FieldArray>,          do_parse!(length: parse_long_int >> array: count!(parse_value, length as usize) >> (array)));
named!(pub parse_timestamp<Timestamp>,             call!(parse_long_long_uint));
named!(pub parse_field_table<FieldTable>,          do_parse!(length: parse_long_uint >> table: flat_map!(take!(length as usize), fold_many0!(complete!(pair!(parse_short_string, parse_value)), FieldTable::new(), |mut acc: FieldTable, (key, value)| {
    acc.insert(key, value);
    acc
})) >> (table)));

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

    use nom::IResult;

    const EMPTY: &'static [u8] = b"";

    #[test]
    fn test_parse_value() {
        assert_eq!(parse_value(&[84, 42, 42, 42, 42, 42,  42,  42,  42]),  IResult::Done(EMPTY, AMQPValue::Timestamp(3038287259199220266)));
        assert_eq!(parse_value(&[83, 0,  0,  0,  4,  116, 101, 115, 116]), IResult::Done(EMPTY, AMQPValue::LongString("test".to_string())));
    }

    #[test]
    fn test_parse_type() {
        assert_eq!(parse_type(&[116]), IResult::Done(EMPTY, AMQPType::Boolean));
        assert_eq!(parse_type(&[102]), IResult::Done(EMPTY, AMQPType::Float));
    }

    #[test]
    fn test_parse_boolean() {
        assert_eq!(parse_boolean(&[0]), IResult::Done(EMPTY, false));
        assert_eq!(parse_boolean(&[1]), IResult::Done(EMPTY, true));
    }

    #[test]
    fn test_parse_short_short_uint() {
        assert_eq!(parse_short_short_uint(&[0]),   IResult::Done(EMPTY, 0));
        assert_eq!(parse_short_short_uint(&[255]), IResult::Done(EMPTY, 255));
    }

    #[test]
    fn test_parse_short_short_int() {
        assert_eq!(parse_short_short_int(&[0]),   IResult::Done(EMPTY, 0));
        assert_eq!(parse_short_short_int(&[255]), IResult::Done(EMPTY, -1));
    }

    #[test]
    fn test_parse_short_uint() {
        assert_eq!(parse_short_uint(&[0,   0]),   IResult::Done(EMPTY, 0));
        assert_eq!(parse_short_uint(&[255, 255]), IResult::Done(EMPTY, 65535));
    }

    #[test]
    fn test_parse_short_int() {
        assert_eq!(parse_short_int(&[0,   0]),   IResult::Done(EMPTY, 0));
        assert_eq!(parse_short_int(&[255, 255]), IResult::Done(EMPTY, -1));
    }

    #[test]
    fn test_parse_long_uint() {
        assert_eq!(parse_long_uint(&[0,   0,   0,   0]),   IResult::Done(EMPTY, 0));
        assert_eq!(parse_long_uint(&[255, 255, 255, 255]), IResult::Done(EMPTY, 4294967295));
    }

    #[test]
    fn test_parse_long_int() {
        assert_eq!(parse_long_int(&[0,   0,   0,   0]),   IResult::Done(EMPTY, 0));
        assert_eq!(parse_long_int(&[255, 255, 255, 255]), IResult::Done(EMPTY, -1));
    }

    #[test]
    fn test_parse_long_long_uint() {
        assert_eq!(parse_long_long_uint(&[0,   0,   0,   0,   0,   0,   0,   0]),   IResult::Done(EMPTY, 0));
        assert_eq!(parse_long_long_uint(&[255, 255, 255, 255, 255, 255, 255, 255]), IResult::Done(EMPTY, 18446744073709551615));
    }

    #[test]
    fn test_parse_long_long_int() {
        assert_eq!(parse_long_long_int(&[0,   0,   0,   0,   0,   0,   0,   0]),   IResult::Done(EMPTY, 0));
        assert_eq!(parse_long_long_int(&[255, 255, 255, 255, 255, 255, 255, 255]), IResult::Done(EMPTY, -1));
    }

    #[test]
    fn test_parse_float() {
        assert_eq!(parse_float(&[0,  0,  0,  0]),  IResult::Done(EMPTY, 0.));
        assert_eq!(parse_float(&[42, 42, 42, 42]), IResult::Done(EMPTY, 0.00000000000015113662));
    }

    #[test]
    fn test_parse_double() {
        assert_eq!(parse_double(&[0,  0,  0,  0,  0,  0,  0,  0]),  IResult::Done(EMPTY, 0.));
        assert_eq!(parse_double(&[42, 42, 42, 42, 42, 42, 42, 42]), IResult::Done(EMPTY, 0.0000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000014260258159703532));
    }

    #[test]
    fn test_parse_decimal_value() {
        assert_eq!(parse_decimal_value(&[0,   0,   0,   0,   0]),   IResult::Done(EMPTY, DecimalValue { scale: 0,   value: 0          }));
        assert_eq!(parse_decimal_value(&[255, 255, 255, 255, 255]), IResult::Done(EMPTY, DecimalValue { scale: 255, value: 4294967295 }));
    }

    #[test]
    fn test_parse_short_string() {
        assert_eq!(parse_short_string(&[0]),                     IResult::Done(EMPTY, ShortString::new()));
        assert_eq!(parse_short_string(&[4, 116, 101, 115, 116]), IResult::Done(EMPTY, "test".to_string()));
    }

    #[test]
    fn test_parse_long_string() {
        assert_eq!(parse_long_string(&[0, 0, 0, 0]),                     IResult::Done(EMPTY, LongString::new()));
        assert_eq!(parse_long_string(&[0, 0, 0, 4, 116, 101, 115, 116]), IResult::Done(EMPTY, "test".to_string()));
    }

    #[test]
    fn test_parse_field_array() {
        assert_eq!(parse_field_array(&[0, 0, 0, 0]),                                 IResult::Done(EMPTY, FieldArray::new()));
        assert_eq!(parse_field_array(&[0, 0, 0, 2, 115, 4, 116, 101, 115, 116, 86]), IResult::Done(EMPTY, vec![AMQPValue::ShortString("test".to_string()), AMQPValue::Void]));
    }

    #[test]
    fn test_parse_timestamp() {
        assert_eq!(parse_timestamp(&[0,   0,   0,   0,   0,   0,   0,   0]),   IResult::Done(EMPTY, 0));
        assert_eq!(parse_timestamp(&[255, 255, 255, 255, 255, 255, 255, 255]), IResult::Done(EMPTY, 18446744073709551615));
    }

    #[test]
    fn test_parse_field_table() {
        let mut table = FieldTable::new();
        table.insert("test".to_string(), AMQPValue::ShortString("test".to_string()));
        table.insert("tt".to_string(),   AMQPValue::Void);
        assert_eq!(parse_field_table(&[0, 0, 0, 0]),                                                                      IResult::Done(EMPTY, FieldTable::new()));
        assert_eq!(parse_field_table(&[0, 0, 0, 15, 4, 116, 101, 115, 116, 115, 4, 116, 101, 115, 116, 2, 116, 116, 86]), IResult::Done(EMPTY, table));
    }
}
<|endoftext|>


<commit_before>use syntax::{
    ast::{self, AstNode, AstToken},
    match_ast, SyntaxElement, TextRange, TextSize, T,
};

use crate::{AssistContext, AssistId, AssistKind, Assists};

// Assist: remove_dbg
//
// Removes `dbg!()` macro call.
//
// ```
// fn main() {
//     $0dbg!(92);
// }
// ```
// ->
// ```
// fn main() {
//     92;
// }
// ```
pub(crate) fn remove_dbg(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let macro_call = ctx.find_node_at_offset::<ast::MacroCall>()?;
    let new_contents = adjusted_macro_contents(&macro_call)?;

    let parent = macro_call.syntax().parent();

    let macro_text_range = if let Some(it) = parent.as_ref() {
        if new_contents.is_empty() {
            match_ast! {
                match it {
                    ast::BlockExpr(_it) => {
                        macro_call.syntax()
                            .prev_sibling_or_token()
                            .and_then(whitespace_start)
                            .map(|start| TextRange::new(start, macro_call.syntax().text_range().end()))
                            .unwrap_or_else(|| macro_call.syntax().text_range())
                    },
                    ast::ExprStmt(it) => {
                        let start = it
                            .syntax()
                            .prev_sibling_or_token()
                            .and_then(whitespace_start)
                            .unwrap_or_else(|| it.syntax().text_range().start());
                        let end = it.syntax().text_range().end();

                        TextRange::new(start, end)
                    },
                    _ => macro_call.syntax().text_range()
                }
            }
        } else {
            macro_call.syntax().text_range()
        }
    } else {
        macro_call.syntax().text_range()
    };

    let macro_end = if macro_call.semicolon_token().is_some() {
        macro_text_range.end() - TextSize::of(';')
    } else {
        macro_text_range.end()
    };

    acc.add(
        AssistId("remove_dbg", AssistKind::Refactor),
        "Remove dbg!()",
        macro_text_range,
        |builder| {
            builder.replace(
                TextRange::new(macro_text_range.start(), macro_end),
                if new_contents.is_empty() && parent.and_then(ast::LetStmt::cast).is_some() {
                    ast::make::expr_unit().to_string()
                } else {
                    new_contents
                },
            );
        },
    )
}

fn whitespace_start(it: SyntaxElement) -> Option<TextSize> {
    Some(it.into_token().and_then(ast::Whitespace::cast)?.syntax().text_range().start())
}

fn adjusted_macro_contents(macro_call: &ast::MacroCall) -> Option<String> {
    let contents = get_valid_macrocall_contents(macro_call, "dbg")?;
    let macro_text_with_brackets = macro_call.token_tree()?.syntax().text();
    let macro_text_in_brackets = macro_text_with_brackets.slice(TextRange::new(
        TextSize::of('('),
        macro_text_with_brackets.len() - TextSize::of(')'),
    ));

    Some(
        if !is_leaf_or_control_flow_expr(macro_call)
            && needs_parentheses_around_macro_contents(contents)
        {
            format!("({})", macro_text_in_brackets)
        } else {
            macro_text_in_brackets.to_string()
        },
    )
}

fn is_leaf_or_control_flow_expr(macro_call: &ast::MacroCall) -> bool {
    macro_call.syntax().next_sibling().is_none()
        || match macro_call.syntax().parent() {
            Some(parent) => match_ast! {
                match parent {
                    ast::Condition(_it) => true,
                    ast::MatchExpr(_it) => true,
                    _ => false,
                }
            },
            None => false,
        }
}

/// Verifies that the given macro_call actually matches the given name
/// and contains proper ending tokens, then returns the contents between the ending tokens
fn get_valid_macrocall_contents(
    macro_call: &ast::MacroCall,
    macro_name: &str,
) -> Option<Vec<SyntaxElement>> {
    let path = macro_call.path()?;
    let name_ref = path.segment()?.name_ref()?;

    // Make sure it is actually a dbg-macro call, dbg followed by !
    let excl = path.syntax().next_sibling_or_token()?;
    if name_ref.text() != macro_name || excl.kind() != T![!] {
        return None;
    }

    let mut children_with_tokens = macro_call.token_tree()?.syntax().children_with_tokens();
    let first_child = children_with_tokens.next()?;
    let mut contents_between_brackets = children_with_tokens.collect::<Vec<_>>();
    let last_child = contents_between_brackets.pop()?;

    match (first_child.kind(), last_child.kind()) {
        (T!['('], T![')']) | (T!['['], T![']']) | (T!['{'], T!['}']) => {
            Some(contents_between_brackets)
        }
        _ => None,
    }
}

fn needs_parentheses_around_macro_contents(macro_contents: Vec<SyntaxElement>) -> bool {
    if macro_contents.len() < 2 {
        return false;
    }
    let mut macro_contents = macro_contents.into_iter().peekable();
    let mut unpaired_brackets_in_contents = Vec::new();
    while let Some(element) = macro_contents.next() {
        match element.kind() {
            T!['('] | T!['['] | T!['{'] => unpaired_brackets_in_contents.push(element),
            T![')'] => {
                if !matches!(unpaired_brackets_in_contents.pop(), Some(correct_bracket) if correct_bracket.kind() == T!['('])
                {
                    return true;
                }
            }
            T![']'] => {
                if !matches!(unpaired_brackets_in_contents.pop(), Some(correct_bracket) if correct_bracket.kind() == T!['['])
                {
                    return true;
                }
            }
            T!['}'] => {
                if !matches!(unpaired_brackets_in_contents.pop(), Some(correct_bracket) if correct_bracket.kind() == T!['{'])
                {
                    return true;
                }
            }
            symbol_kind => {
                let symbol_not_in_bracket = unpaired_brackets_in_contents.is_empty();
                if symbol_not_in_bracket
                    && symbol_kind != T![:] // paths
                    && (symbol_kind != T![.] // field/method access
                        || macro_contents // range expressions consist of two SyntaxKind::Dot in macro invocations
                            .peek()
                            .map(|element| element.kind() == T![.])
                            .unwrap_or(false))
                    && symbol_kind != T![?] // try operator
                    && (symbol_kind.is_punct() || symbol_kind == T![as])
                {
                    return true;
                }
            }
        }
    }
    !unpaired_brackets_in_contents.is_empty()
}

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

    use super::*;

    #[test]
    fn test_remove_dbg() {
        check_assist(remove_dbg, "$0dbg!(1 + 1)", "1 + 1");

        check_assist(remove_dbg, "dbg!$0((1 + 1))", "(1 + 1)");

        check_assist(remove_dbg, "dbg!(1 $0+ 1)", "1 + 1");

        check_assist(remove_dbg, "let _ = $0dbg!(1 + 1)", "let _ = 1 + 1");

        check_assist(
            remove_dbg,
            "
fn foo(n: usize) {
    if let Some(_) = dbg!(n.$0checked_sub(4)) {
        // ...
    }
}
",
            "
fn foo(n: usize) {
    if let Some(_) = n.checked_sub(4) {
        // ...
    }
}
",
        );

        check_assist(remove_dbg, "$0dbg!(Foo::foo_test()).bar()", "Foo::foo_test().bar()");
    }

    #[test]
    fn test_remove_dbg_with_brackets_and_braces() {
        check_assist(remove_dbg, "dbg![$01 + 1]", "1 + 1");
        check_assist(remove_dbg, "dbg!{$01 + 1}", "1 + 1");
    }

    #[test]
    fn test_remove_dbg_not_applicable() {
        check_assist_not_applicable(remove_dbg, "$0vec![1, 2, 3]");
        check_assist_not_applicable(remove_dbg, "$0dbg(5, 6, 7)");
        check_assist_not_applicable(remove_dbg, "$0dbg!(5, 6, 7");
    }

    #[test]
    fn test_remove_dbg_target() {
        check_assist_target(
            remove_dbg,
            "
fn foo(n: usize) {
    if let Some(_) = dbg!(n.$0checked_sub(4)) {
        // ...
    }
}
",
            "dbg!(n.checked_sub(4))",
        );
    }

    #[test]
    fn test_remove_dbg_keep_semicolon() {
        // https://github.com/rust-analyzer/rust-analyzer/issues/5129#issuecomment-651399779
        // not quite though
        // adding a comment at the end of the line makes
        // the ast::MacroCall to include the semicolon at the end
        check_assist(
            remove_dbg,
            r#"let res = $0dbg!(1 * 20); // needless comment"#,
            r#"let res = 1 * 20; // needless comment"#,
        );
    }

    #[test]
    fn remove_dbg_from_non_leaf_simple_expression() {
        check_assist(
            remove_dbg,
            "
fn main() {
    let mut a = 1;
    while dbg!$0(a) < 10000 {
        a += 1;
    }
}
",
            "
fn main() {
    let mut a = 1;
    while a < 10000 {
        a += 1;
    }
}
",
        );
    }

    #[test]
    fn test_remove_dbg_keep_expression() {
        check_assist(
            remove_dbg,
            r#"let res = $0dbg!(a + b).foo();"#,
            r#"let res = (a + b).foo();"#,
        );

        check_assist(remove_dbg, r#"let res = $0dbg!(2 + 2) * 5"#, r#"let res = (2 + 2) * 5"#);
        check_assist(remove_dbg, r#"let res = $0dbg![2 + 2] * 5"#, r#"let res = (2 + 2) * 5"#);
    }

    #[test]
    fn test_remove_dbg_method_chaining() {
        check_assist(
            remove_dbg,
            r#"let res = $0dbg!(foo().bar()).baz();"#,
            r#"let res = foo().bar().baz();"#,
        );
        check_assist(
            remove_dbg,
            r#"let res = $0dbg!(foo.bar()).baz();"#,
            r#"let res = foo.bar().baz();"#,
        );
    }

    #[test]
    fn test_remove_dbg_field_chaining() {
        check_assist(remove_dbg, r#"let res = $0dbg!(foo.bar).baz;"#, r#"let res = foo.bar.baz;"#);
    }

    #[test]
    fn test_remove_dbg_from_inside_fn() {
        check_assist_target(
            remove_dbg,
            r#"
fn square(x: u32) -> u32 {
    x * x
}

fn main() {
    let x = square(dbg$0!(5 + 10));
    println!("{}", x);
}"#,
            "dbg!(5 + 10)",
        );

        check_assist(
            remove_dbg,
            r#"
fn square(x: u32) -> u32 {
    x * x
}

fn main() {
    let x = square(dbg$0!(5 + 10));
    println!("{}", x);
}"#,
            r#"
fn square(x: u32) -> u32 {
    x * x
}

fn main() {
    let x = square(5 + 10);
    println!("{}", x);
}"#,
        );
    }

    #[test]
    fn test_remove_dbg_try_expr() {
        check_assist(
            remove_dbg,
            r#"let res = $0dbg!(result?).foo();"#,
            r#"let res = result?.foo();"#,
        );
    }

    #[test]
    fn test_remove_dbg_await_expr() {
        check_assist(
            remove_dbg,
            r#"let res = $0dbg!(fut.await).foo();"#,
            r#"let res = fut.await.foo();"#,
        );
    }

    #[test]
    fn test_remove_dbg_as_cast() {
        check_assist(
            remove_dbg,
            r#"let res = $0dbg!(3 as usize).foo();"#,
            r#"let res = (3 as usize).foo();"#,
        );
    }

    #[test]
    fn test_remove_dbg_index_expr() {
        check_assist(
            remove_dbg,
            r#"let res = $0dbg!(array[3]).foo();"#,
            r#"let res = array[3].foo();"#,
        );
        check_assist(
            remove_dbg,
            r#"let res = $0dbg!(tuple.3).foo();"#,
            r#"let res = tuple.3.foo();"#,
        );
    }

    #[test]
    fn test_remove_dbg_range_expr() {
        check_assist(
            remove_dbg,
            r#"let res = $0dbg!(foo..bar).foo();"#,
            r#"let res = (foo..bar).foo();"#,
        );
        check_assist(
            remove_dbg,
            r#"let res = $0dbg!(foo..=bar).foo();"#,
            r#"let res = (foo..=bar).foo();"#,
        );
    }

    #[test]
    fn test_remove_dbg_followed_by_block() {
        check_assist(
            remove_dbg,
            r#"fn foo() {
    if $0dbg!(x || y) {}
}"#,
            r#"fn foo() {
    if x || y {}
}"#,
        );
        check_assist(
            remove_dbg,
            r#"fn foo() {
    while let foo = $0dbg!(&x) {}
}"#,
            r#"fn foo() {
    while let foo = &x {}
}"#,
        );
        check_assist(
            remove_dbg,
            r#"fn foo() {
    if let foo = $0dbg!(&x) {}
}"#,
            r#"fn foo() {
    if let foo = &x {}
}"#,
        );
        check_assist(
            remove_dbg,
            r#"fn foo() {
    match $0dbg!(&x) {}
}"#,
            r#"fn foo() {
    match &x {}
}"#,
        );
    }

    #[test]
    fn test_remove_empty_dbg() {
        check_assist(remove_dbg, r#"fn foo() { $0dbg!(); }"#, r#"fn foo() { }"#);
        check_assist(
            remove_dbg,
            r#"
fn foo() {
    $0dbg!();
}
"#,
            r#"
fn foo() {
}
"#,
        );
        check_assist(
            remove_dbg,
            r#"
fn foo() {
    let test = $0dbg!();
}"#,
            r#"
fn foo() {
    let test = ();
}"#,
        );
        check_assist(
            remove_dbg,
            r#"
fn foo() {
    let t = {
        println!("Hello, world");
        $0dbg!()
    };
}"#,
            r#"
fn foo() {
    let t = {
        println!("Hello, world");
    };
}"#,
        );
    }
}
<commit_msg>Parse input expressions in for dbg! in remove_dbg<commit_after>use itertools::Itertools;
use syntax::{
    ast::{self, AstNode, AstToken},
    match_ast, NodeOrToken, SyntaxElement, TextRange, TextSize, T,
};

use crate::{AssistContext, AssistId, AssistKind, Assists};

// Assist: remove_dbg
//
// Removes `dbg!()` macro call.
//
// ```
// fn main() {
//     $0dbg!(92);
// }
// ```
// ->
// ```
// fn main() {
//     92;
// }
// ```
pub(crate) fn remove_dbg(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let macro_call = ctx.find_node_at_offset::<ast::MacroCall>()?;
    let tt = macro_call.token_tree()?;
    let r_delim = NodeOrToken::Token(tt.right_delimiter_token()?);
    if macro_call.path()?.segment()?.name_ref()?.text() != "dbg"
        || macro_call.excl_token().is_none()
    {
        return None;
    }

    let mac_input = tt.syntax().children_with_tokens().skip(1).take_while(|it| *it != r_delim);
    let input_expressions = mac_input.into_iter().group_by(|tok| tok.kind() == T![,]);
    let input_expressions = input_expressions
        .into_iter()
        .filter_map(|(is_sep, group)| (!is_sep).then(|| group))
        .map(|mut tokens| ast::Expr::parse(&tokens.join("")))
        .collect::<Result<Vec<ast::Expr>, _>>()
        .ok()?;

    let parent = macro_call.syntax().parent()?;
    let parent_is_let_stmt = ast::LetStmt::can_cast(parent.kind());
    let (range, text) = match &*input_expressions {
        // dbg!()
        [] => {
            match_ast! {
                match parent {
                    ast::BlockExpr(__) => {
                        let range = macro_call.syntax().text_range();
                        let range = match whitespace_start(macro_call.syntax().prev_sibling_or_token()) {
                            Some(start) => range.cover_offset(start),
                            None => range,
                        };
                        (range, String::new())
                    },
                    ast::ExprStmt(it) => {
                        let range = it.syntax().text_range();
                        let range = match whitespace_start(it.syntax().prev_sibling_or_token()) {
                            Some(start) => range.cover_offset(start),
                            None => range,
                        };
                        (range, String::new())
                    },
                    _ => (macro_call.syntax().text_range(), "()".to_owned())
                }
            }
        }
        // dbg!(expr0)
        [expr] => {
            let wrap = match ast::Expr::cast(parent) {
                Some(parent) => match (expr, parent) {
                    (ast::Expr::CastExpr(_), ast::Expr::CastExpr(_)) => false,
                    (
                        ast::Expr::BoxExpr(_) | ast::Expr::PrefixExpr(_) | ast::Expr::RefExpr(_),
                        ast::Expr::AwaitExpr(_)
                        | ast::Expr::BinExpr(_)
                        | ast::Expr::CallExpr(_)
                        | ast::Expr::CastExpr(_)
                        | ast::Expr::FieldExpr(_)
                        | ast::Expr::IndexExpr(_)
                        | ast::Expr::MethodCallExpr(_)
                        | ast::Expr::RangeExpr(_)
                        | ast::Expr::TryExpr(_),
                    ) => true,
                    (
                        ast::Expr::BinExpr(_) | ast::Expr::CastExpr(_) | ast::Expr::RangeExpr(_),
                        ast::Expr::AwaitExpr(_)
                        | ast::Expr::BinExpr(_)
                        | ast::Expr::CallExpr(_)
                        | ast::Expr::CastExpr(_)
                        | ast::Expr::FieldExpr(_)
                        | ast::Expr::IndexExpr(_)
                        | ast::Expr::MethodCallExpr(_)
                        | ast::Expr::PrefixExpr(_)
                        | ast::Expr::RangeExpr(_)
                        | ast::Expr::RefExpr(_)
                        | ast::Expr::TryExpr(_),
                    ) => true,
                    _ => false,
                },
                None => false,
            };
            (
                macro_call.syntax().text_range(),
                if wrap { format!("({})", expr) } else { expr.to_string() },
            )
        }
        // dbg!(expr0, expr1, ...)
        exprs => (macro_call.syntax().text_range(), format!("({})", exprs.iter().format(","))),
    };

    let range = if macro_call.semicolon_token().is_some() && parent_is_let_stmt {
        TextRange::new(range.start(), range.end() - TextSize::of(";"))
    } else {
        range
    };

    acc.add(AssistId("remove_dbg", AssistKind::Refactor), "Remove dbg!()", range, |builder| {
        builder.replace(range, text);
    })
}

fn whitespace_start(it: Option<SyntaxElement>) -> Option<TextSize> {
    Some(it?.into_token().and_then(ast::Whitespace::cast)?.syntax().text_range().start())
}

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

    use super::*;

    fn check(ra_fixture_before: &str, ra_fixture_after: &str) {
        check_assist(
            remove_dbg,
            &format!("fn main() {{\n{}\n}}", ra_fixture_before),
            &format!("fn main() {{\n{}\n}}", ra_fixture_after),
        );
    }

    #[test]
    fn test_remove_dbg() {
        check("$0dbg!(1 + 1)", "1 + 1");
        check("dbg!$0((1 + 1))", "(1 + 1)");
        check("dbg!(1 $0+ 1)", "1 + 1");
        check("let _ = $0dbg!(1 + 1)", "let _ = 1 + 1");
        check(
            "if let Some(_) = dbg!(n.$0checked_sub(4)) {}",
            "if let Some(_) = n.checked_sub(4) {}",
        );
        check("$0dbg!(Foo::foo_test()).bar()", "Foo::foo_test().bar()");
    }

    #[test]
    fn test_remove_dbg_with_brackets_and_braces() {
        check("dbg![$01 + 1]", "1 + 1");
        check("dbg!{$01 + 1}", "1 + 1");
    }

    #[test]
    fn test_remove_dbg_not_applicable() {
        check_assist_not_applicable(remove_dbg, "fn main() {$0vec![1, 2, 3]}");
        check_assist_not_applicable(remove_dbg, "fn main() {$0dbg(5, 6, 7)}");
        check_assist_not_applicable(remove_dbg, "fn main() {$0dbg!(5, 6, 7}");
    }

    #[test]
    fn test_remove_dbg_keep_semicolon() {
        // https://github.com/rust-analyzer/rust-analyzer/issues/5129#issuecomment-651399779
        // not quite though
        // adding a comment at the end of the line makes
        // the ast::MacroCall to include the semicolon at the end
        check(
            r#"let res = $0dbg!(1 * 20); // needless comment"#,
            r#"let res = 1 * 20; // needless comment"#,
        );
    }

    #[test]
    fn remove_dbg_from_non_leaf_simple_expression() {
        check_assist(
            remove_dbg,
            r"
fn main() {
    let mut a = 1;
    while dbg!$0(a) < 10000 {
        a += 1;
    }
}
",
            r"
fn main() {
    let mut a = 1;
    while a < 10000 {
        a += 1;
    }
}
",
        );
    }

    #[test]
    fn test_remove_dbg_keep_expression() {
        check(r#"let res = $0dbg!(a + b).foo();"#, r#"let res = (a + b).foo();"#);
        check(r#"let res = $0dbg!(2 + 2) * 5"#, r#"let res = (2 + 2) * 5"#);
        check(r#"let res = $0dbg![2 + 2] * 5"#, r#"let res = (2 + 2) * 5"#);
    }

    #[test]
    fn test_remove_dbg_method_chaining() {
        check(r#"let res = $0dbg!(foo().bar()).baz();"#, r#"let res = foo().bar().baz();"#);
        check(r#"let res = $0dbg!(foo.bar()).baz();"#, r#"let res = foo.bar().baz();"#);
    }

    #[test]
    fn test_remove_dbg_field_chaining() {
        check(r#"let res = $0dbg!(foo.bar).baz;"#, r#"let res = foo.bar.baz;"#);
    }

    #[test]
    fn test_remove_dbg_from_inside_fn() {
        check_assist(
            remove_dbg,
            r#"
fn square(x: u32) -> u32 {
    x * x
}

fn main() {
    let x = square(dbg$0!(5 + 10));
    println!("{}", x);
}"#,
            r#"
fn square(x: u32) -> u32 {
    x * x
}

fn main() {
    let x = square(5 + 10);
    println!("{}", x);
}"#,
        );
    }

    #[test]
    fn test_remove_dbg_try_expr() {
        check(r#"let res = $0dbg!(result?).foo();"#, r#"let res = result?.foo();"#);
    }

    #[test]
    fn test_remove_dbg_await_expr() {
        check(r#"let res = $0dbg!(fut.await).foo();"#, r#"let res = fut.await.foo();"#);
    }

    #[test]
    fn test_remove_dbg_as_cast() {
        check(r#"let res = $0dbg!(3 as usize).foo();"#, r#"let res = (3 as usize).foo();"#);
    }

    #[test]
    fn test_remove_dbg_index_expr() {
        check(r#"let res = $0dbg!(array[3]).foo();"#, r#"let res = array[3].foo();"#);
        check(r#"let res = $0dbg!(tuple.3).foo();"#, r#"let res = tuple.3.foo();"#);
    }

    #[test]
    fn test_remove_dbg_range_expr() {
        check(r#"let res = $0dbg!(foo..bar).foo();"#, r#"let res = (foo..bar).foo();"#);
        check(r#"let res = $0dbg!(foo..=bar).foo();"#, r#"let res = (foo..=bar).foo();"#);
    }

    #[test]
    fn test_remove_dbg_followed_by_block() {
        check_assist(
            remove_dbg,
            r#"fn foo() {
    if $0dbg!(x || y) {}
}"#,
            r#"fn foo() {
    if x || y {}
}"#,
        );
        check_assist(
            remove_dbg,
            r#"fn foo() {
    while let foo = $0dbg!(&x) {}
}"#,
            r#"fn foo() {
    while let foo = &x {}
}"#,
        );
        check_assist(
            remove_dbg,
            r#"fn foo() {
    if let foo = $0dbg!(&x) {}
}"#,
            r#"fn foo() {
    if let foo = &x {}
}"#,
        );
        check_assist(
            remove_dbg,
            r#"fn foo() {
    match $0dbg!(&x) {}
}"#,
            r#"fn foo() {
    match &x {}
}"#,
        );
    }

    #[test]
    fn test_remove_empty_dbg() {
        check_assist(remove_dbg, r#"fn foo() { $0dbg!(); }"#, r#"fn foo() { }"#);
        check_assist(
            remove_dbg,
            r#"
fn foo() {
    $0dbg!();
}
"#,
            r#"
fn foo() {
}
"#,
        );
        check_assist(
            remove_dbg,
            r#"
fn foo() {
    let test = $0dbg!();
}"#,
            r#"
fn foo() {
    let test = ();
}"#,
        );
        check_assist(
            remove_dbg,
            r#"
fn foo() {
    let t = {
        println!("Hello, world");
        $0dbg!()
    };
}"#,
            r#"
fn foo() {
    let t = {
        println!("Hello, world");
    };
}"#,
        );
    }
}
<|endoftext|>


<commit_before>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Utils for working with bits

#[cfg(feature = "simd")]
use packed_simd::u8x64;

const BIT_MASK: [u8; 8] = [1, 2, 4, 8, 16, 32, 64, 128];

/// Returns the nearest number that is `>=` than `num` and is a multiple of 64
#[inline]
pub fn round_upto_multiple_of_64(num: usize) -> usize {
    round_upto_power_of_2(num, 64)
}

/// Returns the nearest multiple of `factor` that is `>=` than `num`. Here `factor` must
/// be a power of 2.
pub fn round_upto_power_of_2(num: usize, factor: usize) -> usize {
    debug_assert!(factor > 0 && (factor & (factor - 1)) == 0);
    (num + (factor - 1)) & !(factor - 1)
}

/// Returns whether bit at position `i` in `data` is set or not
#[inline]
pub fn get_bit(data: &[u8], i: usize) -> bool {
    (data[i >> 3] & BIT_MASK[i & 7]) != 0
}

/// Returns whether bit at position `i` in `data` is set or not.
///
/// # Safety
///
/// Note this doesn't do any bound checking, for performance reason. The caller is
/// responsible to guarantee that `i` is within bounds.
#[inline]
pub unsafe fn get_bit_raw(data: *const u8, i: usize) -> bool {
    (*data.add(i >> 3) & BIT_MASK[i & 7]) != 0
}

/// Sets bit at position `i` for `data`
#[inline]
pub fn set_bit(data: &mut [u8], i: usize) {
    data[i >> 3] |= BIT_MASK[i & 7];
}

/// Sets bit at position `i` for `data`
///
/// # Safety
///
/// Note this doesn't do any bound checking, for performance reason. The caller is
/// responsible to guarantee that `i` is within bounds.
#[inline]
pub unsafe fn set_bit_raw(data: *mut u8, i: usize) {
    *data.add(i >> 3) |= BIT_MASK[i & 7];
}

/// Sets bit at position `i` for `data` to 0
#[inline]
pub fn unset_bit(data: &mut [u8], i: usize) {
    data[i >> 3] ^= BIT_MASK[i & 7];
}

/// Sets bit at position `i` for `data` to 0
///
/// # Safety
///
/// Note this doesn't do any bound checking, for performance reason. The caller is
/// responsible to guarantee that `i` is within bounds.
#[inline]
pub unsafe fn unset_bit_raw(data: *mut u8, i: usize) {
    *data.add(i >> 3) ^= BIT_MASK[i & 7];
}

/// Returns the ceil of `value`/`divisor`
#[inline]
pub fn ceil(value: usize, divisor: usize) -> usize {
    let (quot, rem) = (value / divisor, value % divisor);
    if rem > 0 && divisor > 0 {
        quot + 1
    } else {
        quot
    }
}

/// Performs SIMD bitwise binary operations.
///
/// # Safety
///
/// Note that each slice should be 64 bytes and it is the callers responsibility to ensure
/// that this is the case.  If passed slices larger than 64 bytes the operation will only
/// be performed on the first 64 bytes.  Slices less than 64 bytes will panic.
#[cfg(simd)]
pub unsafe fn bitwise_bin_op_simd<F>(left: &[u8], right: &[u8], result: &mut [u8], op: F)
where
    F: Fn(u8x64, u8x64) -> u8x64,
{
    let left_simd = u8x64::from_slice_unaligned_unchecked(left);
    let right_simd = u8x64::from_slice_unaligned_unchecked(right);
    let simd_result = op(left_simd, right_simd);
    simd_result.write_to_slice_unaligned_unchecked(result);
}

#[cfg(test)]
mod tests {
    use std::collections::HashSet;

    use super::*;
    use crate::util::test_util::seedable_rng;
    use rand::Rng;

    #[test]
    fn test_round_upto_multiple_of_64() {
        assert_eq!(0, round_upto_multiple_of_64(0));
        assert_eq!(64, round_upto_multiple_of_64(1));
        assert_eq!(64, round_upto_multiple_of_64(63));
        assert_eq!(64, round_upto_multiple_of_64(64));
        assert_eq!(128, round_upto_multiple_of_64(65));
        assert_eq!(192, round_upto_multiple_of_64(129));
    }

    #[test]
    fn test_get_bit() {
        // 00001101
        assert_eq!(true, get_bit(&[0b00001101], 0));
        assert_eq!(false, get_bit(&[0b00001101], 1));
        assert_eq!(true, get_bit(&[0b00001101], 2));
        assert_eq!(true, get_bit(&[0b00001101], 3));

        // 01001001 01010010
        assert_eq!(true, get_bit(&[0b01001001, 0b01010010], 0));
        assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 1));
        assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 2));
        assert_eq!(true, get_bit(&[0b01001001, 0b01010010], 3));
        assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 4));
        assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 5));
        assert_eq!(true, get_bit(&[0b01001001, 0b01010010], 6));
        assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 7));
        assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 8));
        assert_eq!(true, get_bit(&[0b01001001, 0b01010010], 9));
        assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 10));
        assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 11));
        assert_eq!(true, get_bit(&[0b01001001, 0b01010010], 12));
        assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 13));
        assert_eq!(true, get_bit(&[0b01001001, 0b01010010], 14));
        assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 15));
    }

    #[test]
    fn test_get_bit_raw() {
        const NUM_BYTE: usize = 10;
        let mut buf = vec![0; NUM_BYTE];
        let mut expected = vec![];
        let mut rng = seedable_rng();
        for i in 0..8 * NUM_BYTE {
            let b = rng.gen_bool(0.5);
            expected.push(b);
            if b {
                set_bit(&mut buf[..], i)
            }
        }

        let raw_ptr = buf.as_ptr();
        for (i, b) in expected.iter().enumerate() {
            unsafe {
                assert_eq!(*b, get_bit_raw(raw_ptr, i));
            }
        }
    }

    #[test]
    fn test_set_bit() {
        let mut b = [0b00000000];
        set_bit(&mut b, 0);
        assert_eq!([0b00000001], b);
        set_bit(&mut b, 2);
        assert_eq!([0b00000101], b);
        set_bit(&mut b, 5);
        assert_eq!([0b00100101], b);
    }

    #[test]
    fn test_unset_bit() {
        let mut b = [0b11111111];
        unset_bit(&mut b, 0);
        assert_eq!([0b11111110], b);
        unset_bit(&mut b, 2);
        assert_eq!([0b11111010], b);
        unset_bit(&mut b, 5);
        assert_eq!([0b11011010], b);
    }

    #[test]
    fn test_set_bit_raw() {
        const NUM_BYTE: usize = 10;
        let mut buf = vec![0; NUM_BYTE];
        let mut expected = vec![];
        let mut rng = seedable_rng();
        for i in 0..8 * NUM_BYTE {
            let b = rng.gen_bool(0.5);
            expected.push(b);
            if b {
                unsafe {
                    set_bit_raw(buf.as_mut_ptr(), i);
                }
            }
        }

        let raw_ptr = buf.as_ptr();
        for (i, b) in expected.iter().enumerate() {
            unsafe {
                assert_eq!(*b, get_bit_raw(raw_ptr, i));
            }
        }
    }

    #[test]
    fn test_unset_bit_raw() {
        const NUM_BYTE: usize = 10;
        let mut buf = vec![255; NUM_BYTE];
        let mut expected = vec![];
        let mut rng = seedable_rng();
        for i in 0..8 * NUM_BYTE {
            let b = rng.gen_bool(0.5);
            expected.push(b);
            if !b {
                unsafe {
                    unset_bit_raw(buf.as_mut_ptr(), i);
                }
            }
        }

        let raw_ptr = buf.as_ptr();
        for (i, b) in expected.iter().enumerate() {
            unsafe {
                assert_eq!(*b, get_bit_raw(raw_ptr, i));
            }
        }
    }

    #[test]
    fn test_get_set_bit_roundtrip() {
        const NUM_BYTES: usize = 10;
        const NUM_SETS: usize = 10;

        let mut buffer: [u8; NUM_BYTES * 8] = [0; NUM_BYTES * 8];
        let mut v = HashSet::new();
        let mut rng = seedable_rng();
        for _ in 0..NUM_SETS {
            let offset = rng.gen_range(0, 8 * NUM_BYTES);
            v.insert(offset);
            set_bit(&mut buffer[..], offset);
        }
        for i in 0..NUM_BYTES * 8 {
            assert_eq!(v.contains(&i), get_bit(&buffer[..], i));
        }
    }

    #[test]
    #[cfg(all(any(target_arch = "x86", target_arch = "x86_64")))]
    fn test_ceil() {
        assert_eq!(ceil(0, 1), 0);
        assert_eq!(ceil(1, 1), 1);
        assert_eq!(ceil(1, 2), 1);
        assert_eq!(ceil(1, 8), 1);
        assert_eq!(ceil(7, 8), 1);
        assert_eq!(ceil(8, 8), 1);
        assert_eq!(ceil(9, 8), 2);
        assert_eq!(ceil(9, 9), 1);
        assert_eq!(ceil(10000000000, 10), 1000000000);
        assert_eq!(ceil(10, 10000000000), 1);
        assert_eq!(ceil(10000000000, 1000000000), 10);
    }

    #[test]
    #[cfg(simd)]
    fn test_bitwise_and_simd() {
        let buf1 = [0b00110011u8; 64];
        let buf2 = [0b11110000u8; 64];
        let mut buf3 = [0b00000000; 64];
        unsafe { bitwise_bin_op_simd(&buf1, &buf2, &mut buf3, |a, b| a & b) };
        for i in buf3.iter() {
            assert_eq!(&0b00110000u8, i);
        }
    }

    #[test]
    #[cfg(simd)]
    fn test_bitwise_or_simd() {
        let buf1 = [0b00110011u8; 64];
        let buf2 = [0b11110000u8; 64];
        let mut buf3 = [0b00000000; 64];
        unsafe { bitwise_bin_op_simd(&buf1, &buf2, &mut buf3, |a, b| a | b) };
        for i in buf3.iter() {
            assert_eq!(&0b11110011u8, i);
        }
    }
}
<commit_msg>ARROW-11311: [Rust] Fixed unset_bit<commit_after>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Utils for working with bits

#[cfg(feature = "simd")]
use packed_simd::u8x64;

const BIT_MASK: [u8; 8] = [1, 2, 4, 8, 16, 32, 64, 128];
const UNSET_BIT_MASK: [u8; 8] = [
    255 - 1,
    255 - 2,
    255 - 4,
    255 - 8,
    255 - 16,
    255 - 32,
    255 - 64,
    255 - 128,
];

/// Returns the nearest number that is `>=` than `num` and is a multiple of 64
#[inline]
pub fn round_upto_multiple_of_64(num: usize) -> usize {
    round_upto_power_of_2(num, 64)
}

/// Returns the nearest multiple of `factor` that is `>=` than `num`. Here `factor` must
/// be a power of 2.
pub fn round_upto_power_of_2(num: usize, factor: usize) -> usize {
    debug_assert!(factor > 0 && (factor & (factor - 1)) == 0);
    (num + (factor - 1)) & !(factor - 1)
}

/// Returns whether bit at position `i` in `data` is set or not
#[inline]
pub fn get_bit(data: &[u8], i: usize) -> bool {
    (data[i >> 3] & BIT_MASK[i & 7]) != 0
}

/// Returns whether bit at position `i` in `data` is set or not.
///
/// # Safety
///
/// Note this doesn't do any bound checking, for performance reason. The caller is
/// responsible to guarantee that `i` is within bounds.
#[inline]
pub unsafe fn get_bit_raw(data: *const u8, i: usize) -> bool {
    (*data.add(i >> 3) & BIT_MASK[i & 7]) != 0
}

/// Sets bit at position `i` for `data`
#[inline]
pub fn set_bit(data: &mut [u8], i: usize) {
    data[i >> 3] |= BIT_MASK[i & 7];
}

/// Sets bit at position `i` for `data`
///
/// # Safety
///
/// Note this doesn't do any bound checking, for performance reason. The caller is
/// responsible to guarantee that `i` is within bounds.
#[inline]
pub unsafe fn set_bit_raw(data: *mut u8, i: usize) {
    *data.add(i >> 3) |= BIT_MASK[i & 7];
}

/// Sets bit at position `i` for `data` to 0
#[inline]
pub fn unset_bit(data: &mut [u8], i: usize) {
    data[i >> 3] &= UNSET_BIT_MASK[i & 7];
}

/// Sets bit at position `i` for `data` to 0
///
/// # Safety
///
/// Note this doesn't do any bound checking, for performance reason. The caller is
/// responsible to guarantee that `i` is within bounds.
#[inline]
pub unsafe fn unset_bit_raw(data: *mut u8, i: usize) {
    *data.add(i >> 3) &= UNSET_BIT_MASK[i & 7];
}

/// Returns the ceil of `value`/`divisor`
#[inline]
pub fn ceil(value: usize, divisor: usize) -> usize {
    let (quot, rem) = (value / divisor, value % divisor);
    if rem > 0 && divisor > 0 {
        quot + 1
    } else {
        quot
    }
}

/// Performs SIMD bitwise binary operations.
///
/// # Safety
///
/// Note that each slice should be 64 bytes and it is the callers responsibility to ensure
/// that this is the case.  If passed slices larger than 64 bytes the operation will only
/// be performed on the first 64 bytes.  Slices less than 64 bytes will panic.
#[cfg(simd)]
pub unsafe fn bitwise_bin_op_simd<F>(left: &[u8], right: &[u8], result: &mut [u8], op: F)
where
    F: Fn(u8x64, u8x64) -> u8x64,
{
    let left_simd = u8x64::from_slice_unaligned_unchecked(left);
    let right_simd = u8x64::from_slice_unaligned_unchecked(right);
    let simd_result = op(left_simd, right_simd);
    simd_result.write_to_slice_unaligned_unchecked(result);
}

#[cfg(test)]
mod tests {
    use std::collections::HashSet;

    use super::*;
    use crate::util::test_util::seedable_rng;
    use rand::Rng;

    #[test]
    fn test_round_upto_multiple_of_64() {
        assert_eq!(0, round_upto_multiple_of_64(0));
        assert_eq!(64, round_upto_multiple_of_64(1));
        assert_eq!(64, round_upto_multiple_of_64(63));
        assert_eq!(64, round_upto_multiple_of_64(64));
        assert_eq!(128, round_upto_multiple_of_64(65));
        assert_eq!(192, round_upto_multiple_of_64(129));
    }

    #[test]
    fn test_get_bit() {
        // 00001101
        assert_eq!(true, get_bit(&[0b00001101], 0));
        assert_eq!(false, get_bit(&[0b00001101], 1));
        assert_eq!(true, get_bit(&[0b00001101], 2));
        assert_eq!(true, get_bit(&[0b00001101], 3));

        // 01001001 01010010
        assert_eq!(true, get_bit(&[0b01001001, 0b01010010], 0));
        assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 1));
        assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 2));
        assert_eq!(true, get_bit(&[0b01001001, 0b01010010], 3));
        assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 4));
        assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 5));
        assert_eq!(true, get_bit(&[0b01001001, 0b01010010], 6));
        assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 7));
        assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 8));
        assert_eq!(true, get_bit(&[0b01001001, 0b01010010], 9));
        assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 10));
        assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 11));
        assert_eq!(true, get_bit(&[0b01001001, 0b01010010], 12));
        assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 13));
        assert_eq!(true, get_bit(&[0b01001001, 0b01010010], 14));
        assert_eq!(false, get_bit(&[0b01001001, 0b01010010], 15));
    }

    #[test]
    fn test_get_bit_raw() {
        const NUM_BYTE: usize = 10;
        let mut buf = vec![0; NUM_BYTE];
        let mut expected = vec![];
        let mut rng = seedable_rng();
        for i in 0..8 * NUM_BYTE {
            let b = rng.gen_bool(0.5);
            expected.push(b);
            if b {
                set_bit(&mut buf[..], i)
            }
        }

        let raw_ptr = buf.as_ptr();
        for (i, b) in expected.iter().enumerate() {
            unsafe {
                assert_eq!(*b, get_bit_raw(raw_ptr, i));
            }
        }
    }

    #[test]
    fn test_set_bit() {
        let mut b = [0b00000010];
        set_bit(&mut b, 0);
        assert_eq!([0b00000011], b);
        set_bit(&mut b, 1);
        assert_eq!([0b00000011], b);
        set_bit(&mut b, 7);
        assert_eq!([0b10000011], b);
    }

    #[test]
    fn test_unset_bit() {
        let mut b = [0b11111101];
        unset_bit(&mut b, 0);
        assert_eq!([0b11111100], b);
        unset_bit(&mut b, 1);
        assert_eq!([0b11111100], b);
        unset_bit(&mut b, 7);
        assert_eq!([0b01111100], b);
    }

    #[test]
    fn test_set_bit_raw() {
        const NUM_BYTE: usize = 10;
        let mut buf = vec![0; NUM_BYTE];
        let mut expected = vec![];
        let mut rng = seedable_rng();
        for i in 0..8 * NUM_BYTE {
            let b = rng.gen_bool(0.5);
            expected.push(b);
            if b {
                unsafe {
                    set_bit_raw(buf.as_mut_ptr(), i);
                }
            }
        }

        let raw_ptr = buf.as_ptr();
        for (i, b) in expected.iter().enumerate() {
            unsafe {
                assert_eq!(*b, get_bit_raw(raw_ptr, i));
            }
        }
    }

    #[test]
    fn test_unset_bit_raw() {
        const NUM_BYTE: usize = 10;
        let mut buf = vec![255; NUM_BYTE];
        let mut expected = vec![];
        let mut rng = seedable_rng();
        for i in 0..8 * NUM_BYTE {
            let b = rng.gen_bool(0.5);
            expected.push(b);
            if !b {
                unsafe {
                    unset_bit_raw(buf.as_mut_ptr(), i);
                }
            }
        }

        let raw_ptr = buf.as_ptr();
        for (i, b) in expected.iter().enumerate() {
            unsafe {
                assert_eq!(*b, get_bit_raw(raw_ptr, i));
            }
        }
    }

    #[test]
    fn test_get_set_bit_roundtrip() {
        const NUM_BYTES: usize = 10;
        const NUM_SETS: usize = 10;

        let mut buffer: [u8; NUM_BYTES * 8] = [0; NUM_BYTES * 8];
        let mut v = HashSet::new();
        let mut rng = seedable_rng();
        for _ in 0..NUM_SETS {
            let offset = rng.gen_range(0, 8 * NUM_BYTES);
            v.insert(offset);
            set_bit(&mut buffer[..], offset);
        }
        for i in 0..NUM_BYTES * 8 {
            assert_eq!(v.contains(&i), get_bit(&buffer[..], i));
        }
    }

    #[test]
    #[cfg(all(any(target_arch = "x86", target_arch = "x86_64")))]
    fn test_ceil() {
        assert_eq!(ceil(0, 1), 0);
        assert_eq!(ceil(1, 1), 1);
        assert_eq!(ceil(1, 2), 1);
        assert_eq!(ceil(1, 8), 1);
        assert_eq!(ceil(7, 8), 1);
        assert_eq!(ceil(8, 8), 1);
        assert_eq!(ceil(9, 8), 2);
        assert_eq!(ceil(9, 9), 1);
        assert_eq!(ceil(10000000000, 10), 1000000000);
        assert_eq!(ceil(10, 10000000000), 1);
        assert_eq!(ceil(10000000000, 1000000000), 10);
    }

    #[test]
    #[cfg(simd)]
    fn test_bitwise_and_simd() {
        let buf1 = [0b00110011u8; 64];
        let buf2 = [0b11110000u8; 64];
        let mut buf3 = [0b00000000; 64];
        unsafe { bitwise_bin_op_simd(&buf1, &buf2, &mut buf3, |a, b| a & b) };
        for i in buf3.iter() {
            assert_eq!(&0b00110000u8, i);
        }
    }

    #[test]
    #[cfg(simd)]
    fn test_bitwise_or_simd() {
        let buf1 = [0b00110011u8; 64];
        let buf2 = [0b11110000u8; 64];
        let mut buf3 = [0b00000000; 64];
        unsafe { bitwise_bin_op_simd(&buf1, &buf2, &mut buf3, |a, b| a | b) };
        for i in buf3.iter() {
            assert_eq!(&0b11110011u8, i);
        }
    }
}
<|endoftext|>


<commit_before>use std::unimplemented;

use anyhow::Result;
use log::debug;
use thiserror::Error;
use zydis::enums::Register;

use crate::{
    arch::Arch,
    aspace::AddressSpace,
    module::{Module, Permissions},
    VA,
};

pub mod mmu;
pub mod reg;

#[derive(Error, Debug)]
pub enum EmuError {
    #[error("invalid instruction: {0:#x}")]
    InvalidInstruction(VA),
}

#[derive(Clone)]
pub struct Emulator {
    mem: mmu::MMU,
    reg: reg::Registers,
    dis: zydis::Decoder,
}

impl Emulator {
    pub fn with_arch(arch: Arch) -> Emulator {
        let mut decoder = match arch {
            Arch::X64 => zydis::Decoder::new(zydis::MachineMode::LONG_64, zydis::AddressWidth::_64).unwrap(),
            Arch::X32 => zydis::Decoder::new(zydis::MachineMode::LEGACY_32, zydis::AddressWidth::_32).unwrap(),
        };

        // modes described here: https://github.com/zyantific/zydis/blob/5af06d64432aaa3f6af3cd3e120eefa061b790ab/include/Zydis/Decoder.h#L55
        // start with minimal until we need FULL
        decoder.enable_mode(zydis::DecoderMode::MINIMAL, false).unwrap();

        decoder.enable_mode(zydis::DecoderMode::KNC, false).unwrap();
        decoder.enable_mode(zydis::DecoderMode::MPX, false).unwrap();
        decoder.enable_mode(zydis::DecoderMode::CET, false).unwrap();
        decoder.enable_mode(zydis::DecoderMode::LZCNT, false).unwrap();
        decoder.enable_mode(zydis::DecoderMode::TZCNT, false).unwrap();
        decoder.enable_mode(zydis::DecoderMode::WBNOINVD, false).unwrap();
        decoder.enable_mode(zydis::DecoderMode::CLDEMOTE, false).unwrap();

        Emulator {
            mem: Default::default(),
            reg: Default::default(),
            dis: decoder,
        }
    }

    pub fn from_module(m: &Module) -> Emulator {
        let mut emu = Emulator::with_arch(m.arch);

        for section in m.sections.iter() {
            let mut page_addr = section.virtual_range.start;

            let section_size = section.virtual_range.end - section.virtual_range.start;
            emu.mem
                .mmap(
                    section.virtual_range.start,
                    crate::util::align(section_size, mmu::PAGE_SIZE as u64),
                    Permissions::W,
                )
                .unwrap();

            while page_addr < section.virtual_range.end {
                let mut page = [0u8; mmu::PAGE_SIZE];

                // AddressSpace currently allows non-page-aligned sizes.
                let page_data = if page_addr + mmu::PAGE_SIZE as u64 > section.virtual_range.end {
                    &mut page[..(section.virtual_range.end - page_addr) as usize]
                } else {
                    &mut page[..]
                };

                m.address_space.read_into(page_addr, page_data).unwrap();
                emu.mem.write_page(page_addr, &page[..]).unwrap();
                page_addr += mmu::PAGE_SIZE as u64;
            }

            emu.mem
                .mprotect(
                    section.virtual_range.start,
                    crate::util::align(section_size, mmu::PAGE_SIZE as u64),
                    section.permissions,
                )
                .unwrap();
        }

        emu
    }

    fn fetch(&mut self) -> Result<zydis::DecodedInstruction> {
        // TODO: 32 vs 64
        let pc = self.reg.rip;
        debug!("emu: fetch: {:#x}", pc);

        // TODO: segmentation.

        let mut buf = [0u8; 16];

        // TODO: perms for read should be X, not R.
        self.mem.read(pc, &mut buf[..])?;
        // TODO: if fail, callback.

        let insn = self.dis.decode(&buf[..]);

        if let Ok(Some(insn)) = insn {
            Ok(insn)
        } else {
            Err(EmuError::InvalidInstruction(pc).into())
        }
    }

    fn read_register(&self, reg: Register) -> u64 {
        match reg {
            Register::RAX => self.reg.rax,
            Register::RBX => self.reg.rbx,
            Register::RCX => self.reg.rcx,
            Register::RDX => self.reg.rdx,
            Register::R8 => self.reg.r8,
            Register::R9 => self.reg.r9,
            Register::R10 => self.reg.r10,
            Register::R11 => self.reg.r11,
            Register::R12 => self.reg.r12,
            Register::R13 => self.reg.r13,
            Register::R14 => self.reg.r14,
            Register::R15 => self.reg.r15,
            Register::RSI => self.reg.rsi,
            Register::RDI => self.reg.rdi,
            Register::RSP => self.reg.rsp,
            Register::RBP => self.reg.rbp,
            _ => unimplemented!(),
        }
    }

    fn write_register(&mut self, reg: Register, size: u16, value: u64) {
        let reg = match reg {
            Register::RAX => &mut self.reg.rax,
            Register::RBX => &mut self.reg.rbx,
            Register::RCX => &mut self.reg.rcx,
            Register::RDX => &mut self.reg.rdx,
            Register::R8 => &mut self.reg.r8,
            Register::R9 => &mut self.reg.r9,
            Register::R10 => &mut self.reg.r10,
            Register::R11 => &mut self.reg.r11,
            Register::R12 => &mut self.reg.r12,
            Register::R13 => &mut self.reg.r13,
            Register::R14 => &mut self.reg.r14,
            Register::R15 => &mut self.reg.r15,
            Register::RSI => &mut self.reg.rsi,
            Register::RDI => &mut self.reg.rdi,
            Register::RSP => &mut self.reg.rsp,
            Register::RBP => &mut self.reg.rbp,
            _ => unimplemented!(),
        };

        match size {
            64 => {
                *reg = value;
            }
            _ => unimplemented!(),
        }
    }

    pub fn step(&mut self) -> Result<()> {
        use zydis::enums::{Mnemonic::*, OperandType::*};

        debug!("emu: step: {:#x}", self.reg.rip);
        let insn = self.fetch()?;
        // TODO: handle invalid fetch
        // TODO: handle invalid instruction

        debug!("emu: step: {:#x}: {:#?}", self.reg.rip, insn.mnemonic);
        match insn.mnemonic {
            MOV => {
                //println!("{:#?}", insn);

                let dst = &insn.operands[0];
                let src = &insn.operands[1];

                let value = match src.ty {
                    IMMEDIATE => src.imm.value,
                    REGISTER => self.read_register(src.reg),
                    _ => unimplemented!(),
                };

                match dst.ty {
                    REGISTER => {
                        self.write_register(dst.reg, dst.size, value);
                    }
                    _ => unimplemented!(),
                }

                self.reg.rip += insn.length as u64;
            }
            _ => {
                self.reg.rip += insn.length as u64;
            }
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use crate::{arch::Arch, emu::*, test::*};

    use anyhow::Result;

    const BASE_ADDRESS: u64 = 0x1000;

    #[test]
    fn raw_create() -> Result<()> {
        //init_logging();

        let mut emu: Emulator = Emulator::with_arch(Arch::X64);

        emu.mem.mmap(BASE_ADDRESS, 0x1000, Permissions::RWX)?;

        // 0:  48 c7 c0 01 00 00 00    mov    rax,0x1
        let code = b"\x48\xC7\xC0\x01\x00\x00\x00";
        emu.mem.write(BASE_ADDRESS, &code[..])?;

        emu.reg.rip = BASE_ADDRESS;
        emu.step()?;

        assert_eq!(emu.reg.rip, BASE_ADDRESS + 0x7);
        assert_eq!(emu.reg.rax, 1);

        Ok(())
    }

    #[test]
    fn from_module() -> Result<()> {
        // 0:  48 c7 c0 01 00 00 00    mov    rax,0x1
        let m = load_shellcode64(&b"\x48\xC7\xC0\x01\x00\x00\x00"[..]);

        let mut emu = Emulator::from_module(&m);
        emu.reg.rip = m.address_space.base_address;
        emu.step()?;

        assert_eq!(emu.reg.rip, m.address_space.base_address + 0x7);
        assert_eq!(emu.reg.rax, 1);

        Ok(())
    }

    fn emu_from_sc(code: &[u8]) -> Emulator {
        let m = load_shellcode64(code);
        let mut emu = Emulator::from_module(&m);
        emu.reg.rip = m.address_space.base_address; // 0x1000

        emu.mem.mmap(0x5000, 0x2000, Permissions::RW).unwrap();
        emu.reg.rsp = 0x6000;
        emu.reg.rbp = 0x6000;

        emu
    }

    #[test]
    fn insn_mov_reg_imm() -> Result<()> {
        // 0:  48 c7 c0 01 00 00 00    mov    rax,0x1
        let mut emu = emu_from_sc(&b"\x48\xC7\xC0\x01\x00\x00\x00"[..]);
        emu.step()?;

        assert_eq!(emu.reg.rax, 1);

        Ok(())
    }

    #[test]
    fn insn_mov_reg_reg() -> Result<()> {
        // 0:  48 89 c3                mov    rbx,rax
        let mut emu = emu_from_sc(&b"\x48\x89\xC3"[..]);
        emu.reg.rax = 1;
        emu.step()?;

        assert_eq!(emu.reg.rax, 1);
        assert_eq!(emu.reg.rbx, 1);

        Ok(())
    }
}
<commit_msg>emu: learn to write variants of 64-bit registers<commit_after>use std::unimplemented;

use anyhow::Result;
use log::debug;
use thiserror::Error;
use zydis::enums::Register;

use crate::{
    arch::Arch,
    aspace::AddressSpace,
    module::{Module, Permissions},
    VA,
};

pub mod mmu;
pub mod reg;

#[derive(Error, Debug)]
pub enum EmuError {
    #[error("invalid instruction: {0:#x}")]
    InvalidInstruction(VA),
}

#[derive(Clone)]
pub struct Emulator {
    mem: mmu::MMU,
    reg: reg::Registers,
    dis: zydis::Decoder,
}

impl Emulator {
    pub fn with_arch(arch: Arch) -> Emulator {
        let mut decoder = match arch {
            Arch::X64 => zydis::Decoder::new(zydis::MachineMode::LONG_64, zydis::AddressWidth::_64).unwrap(),
            Arch::X32 => zydis::Decoder::new(zydis::MachineMode::LEGACY_32, zydis::AddressWidth::_32).unwrap(),
        };

        // modes described here: https://github.com/zyantific/zydis/blob/5af06d64432aaa3f6af3cd3e120eefa061b790ab/include/Zydis/Decoder.h#L55
        // start with minimal until we need FULL
        decoder.enable_mode(zydis::DecoderMode::MINIMAL, false).unwrap();

        decoder.enable_mode(zydis::DecoderMode::KNC, false).unwrap();
        decoder.enable_mode(zydis::DecoderMode::MPX, false).unwrap();
        decoder.enable_mode(zydis::DecoderMode::CET, false).unwrap();
        decoder.enable_mode(zydis::DecoderMode::LZCNT, false).unwrap();
        decoder.enable_mode(zydis::DecoderMode::TZCNT, false).unwrap();
        decoder.enable_mode(zydis::DecoderMode::WBNOINVD, false).unwrap();
        decoder.enable_mode(zydis::DecoderMode::CLDEMOTE, false).unwrap();

        Emulator {
            mem: Default::default(),
            reg: Default::default(),
            dis: decoder,
        }
    }

    pub fn from_module(m: &Module) -> Emulator {
        let mut emu = Emulator::with_arch(m.arch);

        for section in m.sections.iter() {
            let mut page_addr = section.virtual_range.start;

            let section_size = section.virtual_range.end - section.virtual_range.start;
            emu.mem
                .mmap(
                    section.virtual_range.start,
                    crate::util::align(section_size, mmu::PAGE_SIZE as u64),
                    Permissions::W,
                )
                .unwrap();

            while page_addr < section.virtual_range.end {
                let mut page = [0u8; mmu::PAGE_SIZE];

                // AddressSpace currently allows non-page-aligned sizes.
                let page_data = if page_addr + mmu::PAGE_SIZE as u64 > section.virtual_range.end {
                    &mut page[..(section.virtual_range.end - page_addr) as usize]
                } else {
                    &mut page[..]
                };

                m.address_space.read_into(page_addr, page_data).unwrap();
                emu.mem.write_page(page_addr, &page[..]).unwrap();
                page_addr += mmu::PAGE_SIZE as u64;
            }

            emu.mem
                .mprotect(
                    section.virtual_range.start,
                    crate::util::align(section_size, mmu::PAGE_SIZE as u64),
                    section.permissions,
                )
                .unwrap();
        }

        emu
    }

    fn fetch(&mut self) -> Result<zydis::DecodedInstruction> {
        // TODO: 32 vs 64
        let pc = self.reg.rip;
        debug!("emu: fetch: {:#x}", pc);

        // TODO: segmentation.

        let mut buf = [0u8; 16];

        // TODO: perms for read should be X, not R.
        self.mem.read(pc, &mut buf[..])?;
        // TODO: if fail, callback.

        let insn = self.dis.decode(&buf[..]);

        if let Ok(Some(insn)) = insn {
            Ok(insn)
        } else {
            Err(EmuError::InvalidInstruction(pc).into())
        }
    }

    fn read_register(&self, reg: Register) -> u64 {
        match reg {
            Register::RAX => self.reg.rax,
            Register::RBX => self.reg.rbx,
            Register::RCX => self.reg.rcx,
            Register::RDX => self.reg.rdx,
            Register::R8 => self.reg.r8,
            Register::R9 => self.reg.r9,
            Register::R10 => self.reg.r10,
            Register::R11 => self.reg.r11,
            Register::R12 => self.reg.r12,
            Register::R13 => self.reg.r13,
            Register::R14 => self.reg.r14,
            Register::R15 => self.reg.r15,
            Register::RSI => self.reg.rsi,
            Register::RDI => self.reg.rdi,
            Register::RSP => self.reg.rsp,
            Register::RBP => self.reg.rbp,
            _ => unimplemented!(),
        }
    }

    fn write_register(&mut self, reg: Register, size: u16, value: u64) {
        let reg = match reg {
            Register::RAX | Register::EAX | Register::AX | Register::AL => &mut self.reg.rax,
            Register::RBX | Register::EBX | Register::BX | Register::BL => &mut self.reg.rbx,
            Register::RCX | Register::ECX | Register::CX | Register::CL => &mut self.reg.rcx,
            Register::RDX | Register::EDX | Register::DX | Register::DL => &mut self.reg.rdx,
            Register::R8 | Register::R8D | Register::R8W | Register::R8B => &mut self.reg.r8,
            Register::R9 | Register::R9D | Register::R9W | Register::R9B => &mut self.reg.r9,
            Register::R10 | Register::R10D | Register::R10W | Register::R10B => &mut self.reg.r10,
            Register::R11 | Register::R11D | Register::R11W | Register::R11B => &mut self.reg.r11,
            Register::R12 | Register::R12D | Register::R12W | Register::R12B => &mut self.reg.r12,
            Register::R13 | Register::R13D | Register::R13W | Register::R13B => &mut self.reg.r13,
            Register::R14 | Register::R14D | Register::R14W | Register::R14B => &mut self.reg.r14,
            Register::R15 | Register::R15D | Register::R15W | Register::R15B => &mut self.reg.r15,
            Register::RSI | Register::ESI | Register::SI | Register::SIL => &mut self.reg.rsi,
            Register::RDI | Register::EDI | Register::DI | Register::DIL => &mut self.reg.rdi,
            Register::RSP | Register::ESP | Register::SP | Register::SPL => &mut self.reg.rsp,
            Register::RBP | Register::EBP | Register::BP | Register::BPL => &mut self.reg.rbp,
            Register::AH => unimplemented!(),
            Register::BH => unimplemented!(),
            Register::CH => unimplemented!(),
            Register::DH => unimplemented!(),
            _ => unimplemented!(),
        };

        match size {
            64 => {
                *reg = value;
            }
            32 => {
                *reg &= 0xFFFF_FFFF_0000_0000;
                *reg |= value & 0xFFFF_FFFF;
            }
            16 => {
                *reg &= 0xFFFF_FFFF_FFFF_0000;
                *reg |= value & 0xFFFF;
            }
            8 => {
                *reg &= 0xFFFF_FFFF_FFFF_FF00;
                *reg |= value & 0xFF;
            }
            _ => unimplemented!(),
        }
    }

    pub fn step(&mut self) -> Result<()> {
        use zydis::enums::{Mnemonic::*, OperandType::*};

        debug!("emu: step: {:#x}", self.reg.rip);
        let insn = self.fetch()?;
        // TODO: handle invalid fetch
        // TODO: handle invalid instruction

        debug!("emu: step: {:#x}: {:#?}", self.reg.rip, insn.mnemonic);
        match insn.mnemonic {
            MOV => {
                //println!("{:#?}", insn);

                let dst = &insn.operands[0];
                let src = &insn.operands[1];

                let value = match src.ty {
                    IMMEDIATE => src.imm.value,
                    REGISTER => self.read_register(src.reg),
                    _ => unimplemented!(),
                };

                match dst.ty {
                    REGISTER => {
                        self.write_register(dst.reg, dst.size, value);
                    }
                    _ => unimplemented!(),
                }

                self.reg.rip += insn.length as u64;
            }
            _ => {
                self.reg.rip += insn.length as u64;
            }
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use crate::{arch::Arch, emu::*, test::*};

    use anyhow::Result;

    const BASE_ADDRESS: u64 = 0x1000;

    #[test]
    fn raw_create() -> Result<()> {
        //init_logging();

        let mut emu: Emulator = Emulator::with_arch(Arch::X64);

        emu.mem.mmap(BASE_ADDRESS, 0x1000, Permissions::RWX)?;

        // 0:  48 c7 c0 01 00 00 00    mov    rax,0x1
        let code = b"\x48\xC7\xC0\x01\x00\x00\x00";
        emu.mem.write(BASE_ADDRESS, &code[..])?;

        emu.reg.rip = BASE_ADDRESS;
        emu.step()?;

        assert_eq!(emu.reg.rip, BASE_ADDRESS + 0x7);
        assert_eq!(emu.reg.rax, 1);

        Ok(())
    }

    #[test]
    fn from_module() -> Result<()> {
        // 0:  48 c7 c0 01 00 00 00    mov    rax,0x1
        let m = load_shellcode64(&b"\x48\xC7\xC0\x01\x00\x00\x00"[..]);

        let mut emu = Emulator::from_module(&m);
        emu.reg.rip = m.address_space.base_address;
        emu.step()?;

        assert_eq!(emu.reg.rip, m.address_space.base_address + 0x7);
        assert_eq!(emu.reg.rax, 1);

        Ok(())
    }

    fn emu_from_sc(code: &[u8]) -> Emulator {
        let m = load_shellcode64(code);
        let mut emu = Emulator::from_module(&m);
        emu.reg.rip = m.address_space.base_address; // 0x1000

        emu.mem.mmap(0x5000, 0x2000, Permissions::RW).unwrap();
        emu.reg.rsp = 0x6000;
        emu.reg.rbp = 0x6000;

        emu
    }

    #[test]
    fn insn_mov_reg_imm() -> Result<()> {
        // 0:  48 c7 c0 01 00 00 00    mov    rax,0x1
        let mut emu = emu_from_sc(&b"\x48\xC7\xC0\x01\x00\x00\x00"[..]);
        emu.step()?;
        assert_eq!(emu.reg.rax, 1);

        // 0:  b8 01 00 00 00          mov    eax,0x1
        let mut emu = emu_from_sc(&b"\xB8\x01\x00\x00\x00"[..]);
        emu.reg.rax = 0xFFFF_FFFF_FFFF_FFFF;
        emu.step()?;
        assert_eq!(emu.reg.eax(), 1);
        assert_eq!(emu.reg.rax(), 0xFFFF_FFFF_0000_0001);

        // 0:  66 b8 01 00             mov    ax,0x1
        let mut emu = emu_from_sc(&b"\x66\xB8\x01\x00"[..]);
        emu.reg.rax = 0xFFFF_FFFF_FFFF_FFFF;
        emu.step()?;
        assert_eq!(emu.reg.ax(), 1);
        assert_eq!(emu.reg.eax(), 0xFFFF_0001);
        assert_eq!(emu.reg.rax(), 0xFFFF_FFFF_FFFF_0001);

        // 0:  b0 01                   mov    al,0x1
        let mut emu = emu_from_sc(&b"\xB0\x01"[..]);
        emu.reg.rax = 0xFFFF_FFFF_FFFF_FFFF;
        emu.step()?;
        assert_eq!(emu.reg.al(), 1);
        assert_eq!(emu.reg.ax(), 0xFF01);
        assert_eq!(emu.reg.eax(), 0xFFFF_FF01);
        assert_eq!(emu.reg.rax(), 0xFFFF_FFFF_FFFF_FF01);

        Ok(())
    }

    #[test]
    fn insn_mov_reg_reg() -> Result<()> {
        // 0:  48 89 c3                mov    rbx,rax
        let mut emu = emu_from_sc(&b"\x48\x89\xC3"[..]);
        emu.reg.rax = 1;
        emu.step()?;

        assert_eq!(emu.reg.rax, 1);
        assert_eq!(emu.reg.rbx, 1);

        Ok(())
    }
}
<|endoftext|>