<commit_before>use arch::init;

/// Interrupt handler function type.
pub type HandlerFunc = unsafe extern "C" fn();

#[derive(Debug, Clone)]
#[repr(C)]
pub struct ExceptionStackFrame {
    pub instruction_pointer: u64,
    pub code_segment: u64,
    pub cpu_flags: u64,
    pub stack_pointer: u64,
    pub stack_segment: u64,
}

#[derive(Debug, Clone)]
pub struct ExceptionInfo {
    pub instruction_pointer: u64,
    pub code_segment: u64,
    pub cpu_flags: u64,
    pub stack_pointer: u64,
    pub stack_segment: u64,
    pub error_code: Option<u64>,
    pub exception_code: u64
}

#[derive(Debug, Clone)]
pub struct Registers {
    pub rax: u64,
    pub rbx: u64,
    pub rcx: u64,
    pub rdx: u64,
    pub rsi: u64,
    pub rdi: u64,
    pub rbp: u64,
    pub r8: u64,
    pub r9: u64,
    pub r10: u64,
    pub r11: u64,
    pub r12: u64,
    pub r13: u64,
    pub r14: u64,
    pub r15: u64,
}

impl Default for Registers {
    fn default() -> Registers {
        Registers {
            rax: 0, rbx: 0, rcx: 0, rdx: 0, rsi: 0, rdi: 0,
            r8: 0, r9: 0, r10: 0, r11: 0, r12: 0, r13: 0, r14: 0, r15: 0,
            rbp: 0,
        }
    }
}

pub static mut RSP_AFTER_SAVING_REGISTERS: u64 = 0;

unsafe extern "C" fn set_kernel_stack(addr: u64) {
    init::set_kernel_stack(addr);
}

pub unsafe fn switch_to_raw(stack_vaddr: u64, code_start: u64, cpu_flags: u64, code_seg: u64, data_seg: u64) {
    switch_to_raw_naked(stack_vaddr, code_start, cpu_flags, code_seg, data_seg);
}

#[naked]
#[inline(never)]
pub unsafe extern "C" fn switch_to_raw_naked(stack_vaddr: u64, code_start: u64, cpu_flags: u64, code_seg: u64, data_seg: u64) {
    asm!("
       /* save registers */
       push rax
       push rbx
       push rcx
       push rdx
       push rbp
       push rsi
       push rdi
       push r8
       push r9
       push r10
       push r11
       push r12
       push r13
       push r14
       push r15
       mov [$1], rsp

       push r8 /* data seg */
       push rdi /* stack vaddr */
       push rdx /* cpu flags */
       push rcx /* code seg */
       push rsi /* code start */

       mov rdi, rsp
       call $0

       mov rax, [$2]
       mov rbx, [$3]
       mov rcx, [$4]
       mov rdx, [$5]
       mov rsi, [$6]
       mov rdi, [$7]
       mov r8, [$8]
       mov r9, [$9]
       mov r10, [$10]
       mov r11, [$11]
       mov r12, [$12]
       mov r13, [$13]
       mov r14, [$14]
       mov r15, [$15]
       mov rbp, [$16]

       iretq
    "
    ::
         "i"(set_kernel_stack as unsafe extern "C" fn(u64)),
         "i"(&RSP_AFTER_SAVING_REGISTERS),

         "i"(&CUR_REGISTERS.rax),
         "i"(&CUR_REGISTERS.rbx),
         "i"(&CUR_REGISTERS.rcx),
         "i"(&CUR_REGISTERS.rdx),
         "i"(&CUR_REGISTERS.rsi),
         "i"(&CUR_REGISTERS.rdi),
         "i"(&CUR_REGISTERS.r8),
         "i"(&CUR_REGISTERS.r9),
         "i"(&CUR_REGISTERS.r10),
         "i"(&CUR_REGISTERS.r11),
         "i"(&CUR_REGISTERS.r12),
         "i"(&CUR_REGISTERS.r13),
         "i"(&CUR_REGISTERS.r14),
         "i"(&CUR_REGISTERS.r15),
         "i"(&CUR_REGISTERS.rbp),

         "{r8}"(data_seg),
         "{rdi}"(stack_vaddr),
         "{rdx}"(cpu_flags),
         "{rcx}"(code_seg),
         "{rsi}"(code_start)
    ::
    "volatile", "intel");
}

static mut CUR_EXCEPTION_STACK_FRAME: Option<ExceptionStackFrame> = None;
static mut CUR_EXCEPTION_ERROR_CODE: Option<u64> = None;
static mut CUR_EXCEPTION_CODE: Option<u64> = None;
pub static mut CUR_REGISTERS: Registers = Registers {
    rax: 0, rbx: 0, rcx: 0, rdx: 0, rsi: 0, rdi: 0,
    r8: 0, r9: 0, r10: 0, r11: 0, r12: 0, r13: 0, r14: 0, r15: 0, rbp: 0
};

pub unsafe fn set_cur_registers(registers: Registers) {
    CUR_REGISTERS = registers;
}

pub unsafe fn cur_registers() -> Registers {
    CUR_REGISTERS.clone()
}

pub unsafe extern "C" fn store_exception_stack(exception_raw: *const ExceptionStackFrame, exception_code: u64) {
    let exception = unsafe {&*exception_raw};
    CUR_EXCEPTION_STACK_FRAME = Some(exception.clone());
    CUR_EXCEPTION_ERROR_CODE = None;
    CUR_EXCEPTION_CODE = Some(exception_code);
}

pub unsafe extern "C" fn store_error_exception_stack(exception_raw: *const ExceptionStackFrame, error_code: u64, exception_code: u64) {
    let exception = unsafe {&*exception_raw};
    CUR_EXCEPTION_STACK_FRAME = Some(exception.clone());
    CUR_EXCEPTION_ERROR_CODE = Some(error_code);
    CUR_EXCEPTION_CODE = Some(exception_code);
}

macro_rules! return_to_raw_fn {
    ($name: ident, $exception_code: expr) => (
        #[naked]
        #[inline(never)]
        pub unsafe extern "C" fn $name() {
            use ::arch::interrupt::switch::{RSP_AFTER_SAVING_REGISTERS, CUR_REGISTERS};

            asm!("mov [$2], rax
                  mov [$3], rbx
                  mov [$4], rcx
                  mov [$5], rdx
                  mov [$6], rsi
                  mov [$7], rdi
                  mov [$8], r8
                  mov [$9], r9
                  mov [$10], r10
                  mov [$11], r11
                  mov [$12], r12
                  mov [$13], r13
                  mov [$14], r14
                  mov [$15], r15
                  mov [$16], rbp

                  mov rdi, rsp
                  sub rsp, 8
                  call $0

                  mov rsp, [$1]
                  pop r15
                  pop r14
                  pop r13
                  pop r12
                  pop r11
                  pop r10
                  pop r9
                  pop r8
                  pop rdi
                  pop rsi
                  pop rbp
                  pop rdx
                  pop rcx
                  pop rbx
                  pop rax"
                 ::

                 "i"(::arch::interrupt::switch::store_exception_stack as unsafe extern "C" fn(*const ::arch::interrupt::switch::ExceptionStackFrame, u64)),
                 "i"(&RSP_AFTER_SAVING_REGISTERS),

                 "i"(&CUR_REGISTERS.rax),
                 "i"(&CUR_REGISTERS.rbx),
                 "i"(&CUR_REGISTERS.rcx),
                 "i"(&CUR_REGISTERS.rdx),
                 "i"(&CUR_REGISTERS.rsi),
                 "i"(&CUR_REGISTERS.rdi),
                 "i"(&CUR_REGISTERS.r8),
                 "i"(&CUR_REGISTERS.r9),
                 "i"(&CUR_REGISTERS.r10),
                 "i"(&CUR_REGISTERS.r11),
                 "i"(&CUR_REGISTERS.r12),
                 "i"(&CUR_REGISTERS.r13),
                 "i"(&CUR_REGISTERS.r14),
                 "i"(&CUR_REGISTERS.r15),
                 "i"(&CUR_REGISTERS.rbp),

                 "{rsi}"($exception_code)
                 :: "volatile", "intel");
        }
    )
}

macro_rules! return_error_to_raw_fn {
    ($name: ident, $exception_code: expr) => (
        #[naked]
        #[inline(never)]
        pub unsafe extern "C" fn $name() {
            use ::arch::interrupt::switch::{RSP_AFTER_SAVING_REGISTERS, CUR_REGISTERS};

            asm!("mov [$2], rax
                  mov [$3], rbx
                  mov [$4], rcx
                  mov [$5], rdx
                  mov [$6], rsi
                  mov [$7], rdi
                  mov [$8], r8
                  mov [$9], r9
                  mov [$10], r10
                  mov [$11], r11
                  mov [$12], r12
                  mov [$13], r13
                  mov [$14], r14
                  mov [$15], r15
                  mov [$16], rbp

                  pop rsi
                  mov rdi, rsp
                  sub rsp, 8
                  call $0

                  mov rsp, [$1]
                  pop r15
                  pop r14
                  pop r13
                  pop r12
                  pop r11
                  pop r10
                  pop r9
                  pop r8
                  pop rdi
                  pop rsi
                  pop rbp
                  pop rdx
                  pop rcx
                  pop rbx
                  pop rax"
                 ::

                 "i"(::arch::interrupt::switch::store_error_exception_stack as unsafe extern "C" fn(*const ::arch::interrupt::switch::ExceptionStackFrame, u64, u64)),
                 "i"(&RSP_AFTER_SAVING_REGISTERS),

                 "i"(&CUR_REGISTERS.rax),
                 "i"(&CUR_REGISTERS.rbx),
                 "i"(&CUR_REGISTERS.rcx),
                 "i"(&CUR_REGISTERS.rdx),
                 "i"(&CUR_REGISTERS.rsi),
                 "i"(&CUR_REGISTERS.rdi),
                 "i"(&CUR_REGISTERS.r8),
                 "i"(&CUR_REGISTERS.r9),
                 "i"(&CUR_REGISTERS.r10),
                 "i"(&CUR_REGISTERS.r11),
                 "i"(&CUR_REGISTERS.r12),
                 "i"(&CUR_REGISTERS.r13),
                 "i"(&CUR_REGISTERS.r14),
                 "i"(&CUR_REGISTERS.r15),
                 "i"(&CUR_REGISTERS.rbp),

                 "{rdx}"($exception_code)
                 : "rdi" : "volatile", "intel");
        }
    )
}

pub fn last_exception_return_value() -> Option<ExceptionInfo> {
    unsafe {
        CUR_EXCEPTION_STACK_FRAME.clone().map(|exp| {
            ExceptionInfo {
                instruction_pointer: exp.instruction_pointer,
                code_segment: exp.code_segment,
                cpu_flags: exp.cpu_flags,
                stack_pointer: exp.stack_pointer,
                stack_segment: exp.stack_segment,
                error_code: CUR_EXCEPTION_ERROR_CODE,
                exception_code: CUR_EXCEPTION_CODE.unwrap()
            }
        })
    }
}
<commit_msg>Fix a bug that caused rsi to contain wrong value in userspace<commit_after>use arch::init;

/// Interrupt handler function type.
pub type HandlerFunc = unsafe extern "C" fn();

#[derive(Debug, Clone)]
#[repr(C)]
pub struct ExceptionStackFrame {
    pub instruction_pointer: u64,
    pub code_segment: u64,
    pub cpu_flags: u64,
    pub stack_pointer: u64,
    pub stack_segment: u64,
}

#[derive(Debug, Clone)]
pub struct ExceptionInfo {
    pub instruction_pointer: u64,
    pub code_segment: u64,
    pub cpu_flags: u64,
    pub stack_pointer: u64,
    pub stack_segment: u64,
    pub error_code: Option<u64>,
    pub exception_code: u64
}

#[derive(Debug, Clone)]
pub struct Registers {
    pub rax: u64,
    pub rbx: u64,
    pub rcx: u64,
    pub rdx: u64,
    pub rsi: u64,
    pub rdi: u64,
    pub rbp: u64,
    pub r8: u64,
    pub r9: u64,
    pub r10: u64,
    pub r11: u64,
    pub r12: u64,
    pub r13: u64,
    pub r14: u64,
    pub r15: u64,
}

impl Default for Registers {
    fn default() -> Registers {
        Registers {
            rax: 0, rbx: 0, rcx: 0, rdx: 0, rsi: 0, rdi: 0,
            r8: 0, r9: 0, r10: 0, r11: 0, r12: 0, r13: 0, r14: 0, r15: 0,
            rbp: 0,
        }
    }
}

pub static mut RSP_AFTER_SAVING_REGISTERS: u64 = 0;

unsafe extern "C" fn set_kernel_stack(addr: u64) {
    init::set_kernel_stack(addr);
}

pub unsafe fn switch_to_raw(stack_vaddr: u64, code_start: u64, cpu_flags: u64, code_seg: u64, data_seg: u64) {
    switch_to_raw_naked(stack_vaddr, code_start, cpu_flags, code_seg, data_seg);
}

#[naked]
#[inline(never)]
pub unsafe extern "C" fn switch_to_raw_naked(stack_vaddr: u64, code_start: u64, cpu_flags: u64, code_seg: u64, data_seg: u64) {
    asm!("
       /* save registers */
       push rax
       push rbx
       push rcx
       push rdx
       push rbp
       push rsi
       push rdi
       push r8
       push r9
       push r10
       push r11
       push r12
       push r13
       push r14
       push r15
       mov [$1], rsp

       push r8 /* data seg */
       push rdi /* stack vaddr */
       push rdx /* cpu flags */
       push rcx /* code seg */
       push rsi /* code start */

       mov rdi, rsp
       call $0

       mov rax, [$2]
       mov rbx, [$3]
       mov rcx, [$4]
       mov rdx, [$5]
       mov rsi, [$6]
       mov rdi, [$7]
       mov r8, [$8]
       mov r9, [$9]
       mov r10, [$10]
       mov r11, [$11]
       mov r12, [$12]
       mov r13, [$13]
       mov r14, [$14]
       mov r15, [$15]
       mov rbp, [$16]

       iretq
    "
    ::
         "i"(set_kernel_stack as unsafe extern "C" fn(u64)),
         "i"(&RSP_AFTER_SAVING_REGISTERS),

         "i"(&CUR_REGISTERS.rax),
         "i"(&CUR_REGISTERS.rbx),
         "i"(&CUR_REGISTERS.rcx),
         "i"(&CUR_REGISTERS.rdx),
         "i"(&CUR_REGISTERS.rsi),
         "i"(&CUR_REGISTERS.rdi),
         "i"(&CUR_REGISTERS.r8),
         "i"(&CUR_REGISTERS.r9),
         "i"(&CUR_REGISTERS.r10),
         "i"(&CUR_REGISTERS.r11),
         "i"(&CUR_REGISTERS.r12),
         "i"(&CUR_REGISTERS.r13),
         "i"(&CUR_REGISTERS.r14),
         "i"(&CUR_REGISTERS.r15),
         "i"(&CUR_REGISTERS.rbp),

         "{r8}"(data_seg),
         "{rdi}"(stack_vaddr),
         "{rdx}"(cpu_flags),
         "{rcx}"(code_seg),
         "{rsi}"(code_start)
    ::
    "volatile", "intel");
}

static mut CUR_EXCEPTION_STACK_FRAME: Option<ExceptionStackFrame> = None;
static mut CUR_EXCEPTION_ERROR_CODE: Option<u64> = None;
static mut CUR_EXCEPTION_CODE: Option<u64> = None;
pub static mut CUR_REGISTERS: Registers = Registers {
    rax: 0, rbx: 0, rcx: 0, rdx: 0, rsi: 0, rdi: 0,
    r8: 0, r9: 0, r10: 0, r11: 0, r12: 0, r13: 0, r14: 0, r15: 0, rbp: 0
};

pub unsafe fn set_cur_registers(registers: Registers) {
    CUR_REGISTERS = registers;
}

pub unsafe fn cur_registers() -> Registers {
    CUR_REGISTERS.clone()
}

pub unsafe extern "C" fn store_exception_stack(exception_raw: *const ExceptionStackFrame, exception_code: u64) {
    let exception = unsafe {&*exception_raw};
    CUR_EXCEPTION_STACK_FRAME = Some(exception.clone());
    CUR_EXCEPTION_ERROR_CODE = None;
    CUR_EXCEPTION_CODE = Some(exception_code);
}

pub unsafe extern "C" fn store_error_exception_stack(exception_raw: *const ExceptionStackFrame, error_code: u64, exception_code: u64) {
    let exception = unsafe {&*exception_raw};
    CUR_EXCEPTION_STACK_FRAME = Some(exception.clone());
    CUR_EXCEPTION_ERROR_CODE = Some(error_code);
    CUR_EXCEPTION_CODE = Some(exception_code);
}

macro_rules! return_to_raw_fn {
    ($name: ident, $exception_code: expr) => (
        #[naked]
        #[inline(never)]
        pub unsafe extern "C" fn $name() {
            use ::arch::interrupt::switch::{RSP_AFTER_SAVING_REGISTERS, CUR_REGISTERS};

            asm!("mov [$2], rax
                  mov [$3], rbx
                  mov [$4], rcx
                  mov [$5], rdx
                  mov [$6], rsi
                  mov [$7], rdi
                  mov [$8], r8
                  mov [$9], r9
                  mov [$10], r10
                  mov [$11], r11
                  mov [$12], r12
                  mov [$13], r13
                  mov [$14], r14
                  mov [$15], r15
                  mov [$16], rbp

                  mov rsi, $17
                  mov rdi, rsp
                  sub rsp, 8
                  call $0

                  mov rsp, [$1]
                  pop r15
                  pop r14
                  pop r13
                  pop r12
                  pop r11
                  pop r10
                  pop r9
                  pop r8
                  pop rdi
                  pop rsi
                  pop rbp
                  pop rdx
                  pop rcx
                  pop rbx
                  pop rax"
                 ::

                 "i"(::arch::interrupt::switch::store_exception_stack as unsafe extern "C" fn(*const ::arch::interrupt::switch::ExceptionStackFrame, u64)),
                 "i"(&RSP_AFTER_SAVING_REGISTERS),

                 "i"(&CUR_REGISTERS.rax),
                 "i"(&CUR_REGISTERS.rbx),
                 "i"(&CUR_REGISTERS.rcx),
                 "i"(&CUR_REGISTERS.rdx),
                 "i"(&CUR_REGISTERS.rsi),
                 "i"(&CUR_REGISTERS.rdi),
                 "i"(&CUR_REGISTERS.r8),
                 "i"(&CUR_REGISTERS.r9),
                 "i"(&CUR_REGISTERS.r10),
                 "i"(&CUR_REGISTERS.r11),
                 "i"(&CUR_REGISTERS.r12),
                 "i"(&CUR_REGISTERS.r13),
                 "i"(&CUR_REGISTERS.r14),
                 "i"(&CUR_REGISTERS.r15),
                 "i"(&CUR_REGISTERS.rbp),

                 "i"($exception_code)
                 :: "volatile", "intel");
        }
    )
}

macro_rules! return_error_to_raw_fn {
    ($name: ident, $exception_code: expr) => (
        #[naked]
        #[inline(never)]
        pub unsafe extern "C" fn $name() {
            use ::arch::interrupt::switch::{RSP_AFTER_SAVING_REGISTERS, CUR_REGISTERS};

            asm!("mov [$2], rax
                  mov [$3], rbx
                  mov [$4], rcx
                  mov [$5], rdx
                  mov [$6], rsi
                  mov [$7], rdi
                  mov [$8], r8
                  mov [$9], r9
                  mov [$10], r10
                  mov [$11], r11
                  mov [$12], r12
                  mov [$13], r13
                  mov [$14], r14
                  mov [$15], r15
                  mov [$16], rbp

                  mov rdx, $17
                  pop rsi
                  mov rdi, rsp
                  sub rsp, 8
                  call $0

                  mov rsp, [$1]
                  pop r15
                  pop r14
                  pop r13
                  pop r12
                  pop r11
                  pop r10
                  pop r9
                  pop r8
                  pop rdi
                  pop rsi
                  pop rbp
                  pop rdx
                  pop rcx
                  pop rbx
                  pop rax"
                 ::

                 "i"(::arch::interrupt::switch::store_error_exception_stack as unsafe extern "C" fn(*const ::arch::interrupt::switch::ExceptionStackFrame, u64, u64)),
                 "i"(&RSP_AFTER_SAVING_REGISTERS),

                 "i"(&CUR_REGISTERS.rax),
                 "i"(&CUR_REGISTERS.rbx),
                 "i"(&CUR_REGISTERS.rcx),
                 "i"(&CUR_REGISTERS.rdx),
                 "i"(&CUR_REGISTERS.rsi),
                 "i"(&CUR_REGISTERS.rdi),
                 "i"(&CUR_REGISTERS.r8),
                 "i"(&CUR_REGISTERS.r9),
                 "i"(&CUR_REGISTERS.r10),
                 "i"(&CUR_REGISTERS.r11),
                 "i"(&CUR_REGISTERS.r12),
                 "i"(&CUR_REGISTERS.r13),
                 "i"(&CUR_REGISTERS.r14),
                 "i"(&CUR_REGISTERS.r15),
                 "i"(&CUR_REGISTERS.rbp),

                 "i"($exception_code)
                 : "rdi" : "volatile", "intel");
        }
    )
}

pub fn last_exception_return_value() -> Option<ExceptionInfo> {
    unsafe {
        CUR_EXCEPTION_STACK_FRAME.clone().map(|exp| {
            ExceptionInfo {
                instruction_pointer: exp.instruction_pointer,
                code_segment: exp.code_segment,
                cpu_flags: exp.cpu_flags,
                stack_pointer: exp.stack_pointer,
                stack_segment: exp.stack_segment,
                error_code: CUR_EXCEPTION_ERROR_CODE,
                exception_code: CUR_EXCEPTION_CODE.unwrap()
            }
        })
    }
}
<|endoftext|>


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

<%namespace name="helpers" file="/helpers.mako.rs" />

${helpers.four_sides_shorthand("margin", "margin-%s", "specified::LengthPercentageOrAuto::parse",
                               spec="https://drafts.csswg.org/css-box/#propdef-margin",
                               allowed_in_page_rule=True,
                               allow_quirks=True)}
<commit_msg>style: Implement the margin-block/inline shorthands.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at https://mozilla.org/MPL/2.0/. */

<%namespace name="helpers" file="/helpers.mako.rs" />

${helpers.four_sides_shorthand("margin", "margin-%s", "specified::LengthPercentageOrAuto::parse",
                               spec="https://drafts.csswg.org/css-box/#propdef-margin",
                               allowed_in_page_rule=True,
                               allow_quirks=True)}

% for axis in ["block", "inline"]:
    <%
        spec = "https://drafts.csswg.org/css-logical/#propdef-margin-%s" % axis
    %>
    <%helpers:shorthand
        name="margin-${axis}"
        sub_properties="${' '.join(
            'margin-%s-%s' % (axis, side)
            for side in ['start', 'end']
        )}"
        spec="${spec}">

        use crate::parser::Parse;
        use crate::values::specified::length::LengthPercentageOrAuto;
        pub fn parse_value<'i, 't>(
            context: &ParserContext,
            input: &mut Parser<'i, 't>,
        ) -> Result<Longhands, ParseError<'i>> {
            let start_value = LengthPercentageOrAuto::parse(context, input)?;
            let end_value =
                input.try(|input| LengthPercentageOrAuto::parse(context, input)).unwrap_or_else(|_| start_value.clone());

            Ok(expanded! {
                margin_${axis}_start: start_value,
                margin_${axis}_end: end_value,
            })
        }

        impl<'a> ToCss for LonghandsToSerialize<'a>  {
            fn to_css<W>(&self, dest: &mut CssWriter<W>) -> fmt::Result where W: fmt::Write {
                self.margin_${axis}_start.to_css(dest)?;

                if self.margin_${axis}_end != self.margin_${axis}_start {
                    dest.write_str(" ")?;
                    self.margin_${axis}_end.to_css(dest)?;
                }

                Ok(())
            }
        }
    </%helpers:shorthand>
% endfor
<|endoftext|>


<commit_before>use syntax::{
    ast::{self, edit::AstNodeEdit, make},
    AstNode,
};

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

// Assist: pull_assignment_up
//
// Extracts variable assignment to outside an if or match statement.
//
// ```
// fn main() {
//     let mut foo = 6;
//
//     if true {
//         $0foo = 5;
//     } else {
//         foo = 4;
//     }
// }
// ```
// ->
// ```
// fn main() {
//     let mut foo = 6;
//
//     foo = if true {
//         5
//     } else {
//         4
//     };
// }
// ```
pub(crate) fn pull_assignment_up(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let assign_expr = ctx.find_node_at_offset::<ast::BinExpr>()?;
    let name_expr = if assign_expr.op_kind()? == ast::BinOp::Assignment {
        assign_expr.lhs()?
    } else {
        return None;
    };

    let (old_stmt, new_stmt) = if let Some(if_expr) = ctx.find_node_at_offset::<ast::IfExpr>() {
        (
            ast::Expr::cast(if_expr.syntax().to_owned())?,
            exprify_if(&if_expr, &ctx.sema, &name_expr)?.indent(if_expr.indent_level()),
        )
    } else if let Some(match_expr) = ctx.find_node_at_offset::<ast::MatchExpr>() {
        (
            ast::Expr::cast(match_expr.syntax().to_owned())?,
            exprify_match(&match_expr, &ctx.sema, &name_expr)?,
        )
    } else {
        return None;
    };

    let expr_stmt = make::expr_stmt(new_stmt);

    acc.add(
        AssistId("pull_assignment_up", AssistKind::RefactorExtract),
        "Pull assignment up",
        old_stmt.syntax().text_range(),
        move |edit| {
            edit.replace(old_stmt.syntax().text_range(), format!("{} = {};", name_expr, expr_stmt));
        },
    )
}

fn exprify_match(
    match_expr: &ast::MatchExpr,
    sema: &hir::Semantics<ide_db::RootDatabase>,
    name: &ast::Expr,
) -> Option<ast::Expr> {
    let new_arm_list = match_expr
        .match_arm_list()?
        .arms()
        .map(|arm| {
            if let ast::Expr::BlockExpr(block) = arm.expr()? {
                let new_block = exprify_block(&block, sema, name)?.indent(block.indent_level());
                Some(arm.replace_descendant(block, new_block))
            } else {
                None
            }
        })
        .collect::<Option<Vec<_>>>()?;
    let new_arm_list = match_expr
        .match_arm_list()?
        .replace_descendants(match_expr.match_arm_list()?.arms().zip(new_arm_list));
    Some(make::expr_match(match_expr.expr()?, new_arm_list))
}

fn exprify_if(
    statement: &ast::IfExpr,
    sema: &hir::Semantics<ide_db::RootDatabase>,
    name: &ast::Expr,
) -> Option<ast::Expr> {
    let then_branch = exprify_block(&statement.then_branch()?, sema, name)?;
    let else_branch = match statement.else_branch()? {
        ast::ElseBranch::Block(ref block) => {
            ast::ElseBranch::Block(exprify_block(block, sema, name)?)
        }
        ast::ElseBranch::IfExpr(expr) => {
            cov_mark::hit!(test_pull_assignment_up_chained_if);
            ast::ElseBranch::IfExpr(ast::IfExpr::cast(
                exprify_if(&expr, sema, name)?.syntax().to_owned(),
            )?)
        }
    };
    Some(make::expr_if(statement.condition()?, then_branch, Some(else_branch)))
}

fn exprify_block(
    block: &ast::BlockExpr,
    sema: &hir::Semantics<ide_db::RootDatabase>,
    name: &ast::Expr,
) -> Option<ast::BlockExpr> {
    if block.tail_expr().is_some() {
        return None;
    }

    let mut stmts: Vec<_> = block.statements().collect();
    let stmt = stmts.pop()?;

    if let ast::Stmt::ExprStmt(stmt) = stmt {
        if let ast::Expr::BinExpr(expr) = stmt.expr()? {
            if expr.op_kind()? == ast::BinOp::Assignment && is_equivalent(sema, &expr.lhs()?, name)
            {
                // The last statement in the block is an assignment to the name we want
                return Some(make::block_expr(stmts, Some(expr.rhs()?)));
            }
        }
    }
    None
}

fn is_equivalent(
    sema: &hir::Semantics<ide_db::RootDatabase>,
    expr0: &ast::Expr,
    expr1: &ast::Expr,
) -> bool {
    match (expr0, expr1) {
        (ast::Expr::FieldExpr(field_expr0), ast::Expr::FieldExpr(field_expr1)) => {
            cov_mark::hit!(test_pull_assignment_up_field_assignment);
            sema.resolve_field(field_expr0) == sema.resolve_field(field_expr1)
        }
        (ast::Expr::PathExpr(path0), ast::Expr::PathExpr(path1)) => {
            let path0 = path0.path();
            let path1 = path1.path();
            if let (Some(path0), Some(path1)) = (path0, path1) {
                sema.resolve_path(&path0) == sema.resolve_path(&path1)
            } else {
                false
            }
        }
        (ast::Expr::PrefixExpr(prefix0), ast::Expr::PrefixExpr(prefix1))
            if prefix0.op_kind() == Some(ast::PrefixOp::Deref)
                && prefix1.op_kind() == Some(ast::PrefixOp::Deref) =>
        {
            cov_mark::hit!(test_pull_assignment_up_deref);
            if let (Some(prefix0), Some(prefix1)) = (prefix0.expr(), prefix1.expr()) {
                is_equivalent(sema, &prefix0, &prefix1)
            } else {
                false
            }
        }
        _ => false,
    }
}

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

    use crate::tests::{check_assist, check_assist_not_applicable};

    #[test]
    fn test_pull_assignment_up_if() {
        check_assist(
            pull_assignment_up,
            r#"
fn foo() {
    let mut a = 1;

    if true {
        $0a = 2;
    } else {
        a = 3;
    }
}"#,
            r#"
fn foo() {
    let mut a = 1;

    a = if true {
        2
    } else {
        3
    };
}"#,
        );
    }

    #[test]
    fn test_pull_assignment_up_match() {
        check_assist(
            pull_assignment_up,
            r#"
fn foo() {
    let mut a = 1;

    match 1 {
        1 => {
            $0a = 2;
        },
        2 => {
            a = 3;
        },
        3 => {
            a = 4;
        }
    }
}"#,
            r#"
fn foo() {
    let mut a = 1;

    a = match 1 {
        1 => {
            2
        },
        2 => {
            3
        },
        3 => {
            4
        }
    };
}"#,
        );
    }

    #[test]
    fn test_pull_assignment_up_not_last_not_applicable() {
        check_assist_not_applicable(
            pull_assignment_up,
            r#"
fn foo() {
    let mut a = 1;

    if true {
        $0a = 2;
        b = a;
    } else {
        a = 3;
    }
}"#,
        )
    }

    #[test]
    fn test_pull_assignment_up_chained_if() {
        cov_mark::check!(test_pull_assignment_up_chained_if);
        check_assist(
            pull_assignment_up,
            r#"
fn foo() {
    let mut a = 1;

    if true {
        $0a = 2;
    } else if false {
        a = 3;
    } else {
        a = 4;
    }
}"#,
            r#"
fn foo() {
    let mut a = 1;

    a = if true {
        2
    } else if false {
        3
    } else {
        4
    };
}"#,
        );
    }

    #[test]
    fn test_pull_assignment_up_retains_stmts() {
        check_assist(
            pull_assignment_up,
            r#"
fn foo() {
    let mut a = 1;

    if true {
        let b = 2;
        $0a = 2;
    } else {
        let b = 3;
        a = 3;
    }
}"#,
            r#"
fn foo() {
    let mut a = 1;

    a = if true {
        let b = 2;
        2
    } else {
        let b = 3;
        3
    };
}"#,
        )
    }

    #[test]
    fn pull_assignment_up_let_stmt_not_applicable() {
        check_assist_not_applicable(
            pull_assignment_up,
            r#"
fn foo() {
    let mut a = 1;

    let b = if true {
        $0a = 2
    } else {
        a = 3
    };
}"#,
        )
    }

    #[test]
    fn pull_assignment_up_if_missing_assigment_not_applicable() {
        check_assist_not_applicable(
            pull_assignment_up,
            r#"
fn foo() {
    let mut a = 1;

    if true {
        $0a = 2;
    } else {}
}"#,
        )
    }

    #[test]
    fn pull_assignment_up_match_missing_assigment_not_applicable() {
        check_assist_not_applicable(
            pull_assignment_up,
            r#"
fn foo() {
    let mut a = 1;

    match 1 {
        1 => {
            $0a = 2;
        },
        2 => {
            a = 3;
        },
        3 => {},
    }
}"#,
        )
    }

    #[test]
    fn test_pull_assignment_up_field_assignment() {
        cov_mark::check!(test_pull_assignment_up_field_assignment);
        check_assist(
            pull_assignment_up,
            r#"
struct A(usize);

fn foo() {
    let mut a = A(1);

    if true {
        $0a.0 = 2;
    } else {
        a.0 = 3;
    }
}"#,
            r#"
struct A(usize);

fn foo() {
    let mut a = A(1);

    a.0 = if true {
        2
    } else {
        3
    };
}"#,
        )
    }

    #[test]
    fn test_pull_assignment_up_deref() {
        cov_mark::check!(test_pull_assignment_up_deref);
        check_assist(
            pull_assignment_up,
            r#"
fn foo() {
    let mut a = 1;
    let b = &mut a;

    if true {
        $0*b = 2;
    } else {
        *b = 3;
    }
}
"#,
            r#"
fn foo() {
    let mut a = 1;
    let b = &mut a;

    *b = if true {
        2
    } else {
        3
    };
}
"#,
        )
    }
}
<commit_msg>minor: add missing test<commit_after>use syntax::{
    ast::{self, edit::AstNodeEdit, make},
    AstNode,
};

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

// Assist: pull_assignment_up
//
// Extracts variable assignment to outside an if or match statement.
//
// ```
// fn main() {
//     let mut foo = 6;
//
//     if true {
//         $0foo = 5;
//     } else {
//         foo = 4;
//     }
// }
// ```
// ->
// ```
// fn main() {
//     let mut foo = 6;
//
//     foo = if true {
//         5
//     } else {
//         4
//     };
// }
// ```
pub(crate) fn pull_assignment_up(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
    let assign_expr = ctx.find_node_at_offset::<ast::BinExpr>()?;

    let op_kind = assign_expr.op_kind()?;
    if op_kind != ast::BinOp::Assignment {
        cov_mark::hit!(test_cant_pull_non_assignments);
        return None;
    }

    let name_expr = assign_expr.lhs()?;

    let old_stmt: ast::Expr;
    let new_stmt: ast::Expr;

    if let Some(if_expr) = ctx.find_node_at_offset::<ast::IfExpr>() {
        new_stmt = exprify_if(&if_expr, &ctx.sema, &name_expr)?.indent(if_expr.indent_level());
        old_stmt = if_expr.into();
    } else if let Some(match_expr) = ctx.find_node_at_offset::<ast::MatchExpr>() {
        new_stmt = exprify_match(&match_expr, &ctx.sema, &name_expr)?;
        old_stmt = match_expr.into()
    } else {
        return None;
    };

    let expr_stmt = make::expr_stmt(new_stmt);

    acc.add(
        AssistId("pull_assignment_up", AssistKind::RefactorExtract),
        "Pull assignment up",
        old_stmt.syntax().text_range(),
        move |edit| {
            edit.replace(old_stmt.syntax().text_range(), format!("{} = {};", name_expr, expr_stmt));
        },
    )
}

fn exprify_match(
    match_expr: &ast::MatchExpr,
    sema: &hir::Semantics<ide_db::RootDatabase>,
    name: &ast::Expr,
) -> Option<ast::Expr> {
    let new_arm_list = match_expr
        .match_arm_list()?
        .arms()
        .map(|arm| {
            if let ast::Expr::BlockExpr(block) = arm.expr()? {
                let new_block = exprify_block(&block, sema, name)?.indent(block.indent_level());
                Some(arm.replace_descendant(block, new_block))
            } else {
                None
            }
        })
        .collect::<Option<Vec<_>>>()?;
    let new_arm_list = match_expr
        .match_arm_list()?
        .replace_descendants(match_expr.match_arm_list()?.arms().zip(new_arm_list));
    Some(make::expr_match(match_expr.expr()?, new_arm_list))
}

fn exprify_if(
    statement: &ast::IfExpr,
    sema: &hir::Semantics<ide_db::RootDatabase>,
    name: &ast::Expr,
) -> Option<ast::Expr> {
    let then_branch = exprify_block(&statement.then_branch()?, sema, name)?;
    let else_branch = match statement.else_branch()? {
        ast::ElseBranch::Block(block) => ast::ElseBranch::Block(exprify_block(&block, sema, name)?),
        ast::ElseBranch::IfExpr(expr) => {
            cov_mark::hit!(test_pull_assignment_up_chained_if);
            ast::ElseBranch::IfExpr(ast::IfExpr::cast(
                exprify_if(&expr, sema, name)?.syntax().to_owned(),
            )?)
        }
    };
    Some(make::expr_if(statement.condition()?, then_branch, Some(else_branch)))
}

fn exprify_block(
    block: &ast::BlockExpr,
    sema: &hir::Semantics<ide_db::RootDatabase>,
    name: &ast::Expr,
) -> Option<ast::BlockExpr> {
    if block.tail_expr().is_some() {
        return None;
    }

    let mut stmts: Vec<_> = block.statements().collect();
    let stmt = stmts.pop()?;

    if let ast::Stmt::ExprStmt(stmt) = stmt {
        if let ast::Expr::BinExpr(expr) = stmt.expr()? {
            if expr.op_kind()? == ast::BinOp::Assignment && is_equivalent(sema, &expr.lhs()?, name)
            {
                // The last statement in the block is an assignment to the name we want
                return Some(make::block_expr(stmts, Some(expr.rhs()?)));
            }
        }
    }
    None
}

fn is_equivalent(
    sema: &hir::Semantics<ide_db::RootDatabase>,
    expr0: &ast::Expr,
    expr1: &ast::Expr,
) -> bool {
    match (expr0, expr1) {
        (ast::Expr::FieldExpr(field_expr0), ast::Expr::FieldExpr(field_expr1)) => {
            cov_mark::hit!(test_pull_assignment_up_field_assignment);
            sema.resolve_field(field_expr0) == sema.resolve_field(field_expr1)
        }
        (ast::Expr::PathExpr(path0), ast::Expr::PathExpr(path1)) => {
            let path0 = path0.path();
            let path1 = path1.path();
            if let (Some(path0), Some(path1)) = (path0, path1) {
                sema.resolve_path(&path0) == sema.resolve_path(&path1)
            } else {
                false
            }
        }
        (ast::Expr::PrefixExpr(prefix0), ast::Expr::PrefixExpr(prefix1))
            if prefix0.op_kind() == Some(ast::PrefixOp::Deref)
                && prefix1.op_kind() == Some(ast::PrefixOp::Deref) =>
        {
            cov_mark::hit!(test_pull_assignment_up_deref);
            if let (Some(prefix0), Some(prefix1)) = (prefix0.expr(), prefix1.expr()) {
                is_equivalent(sema, &prefix0, &prefix1)
            } else {
                false
            }
        }
        _ => false,
    }
}

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

    use crate::tests::{check_assist, check_assist_not_applicable};

    #[test]
    fn test_pull_assignment_up_if() {
        check_assist(
            pull_assignment_up,
            r#"
fn foo() {
    let mut a = 1;

    if true {
        $0a = 2;
    } else {
        a = 3;
    }
}"#,
            r#"
fn foo() {
    let mut a = 1;

    a = if true {
        2
    } else {
        3
    };
}"#,
        );
    }

    #[test]
    fn test_pull_assignment_up_match() {
        check_assist(
            pull_assignment_up,
            r#"
fn foo() {
    let mut a = 1;

    match 1 {
        1 => {
            $0a = 2;
        },
        2 => {
            a = 3;
        },
        3 => {
            a = 4;
        }
    }
}"#,
            r#"
fn foo() {
    let mut a = 1;

    a = match 1 {
        1 => {
            2
        },
        2 => {
            3
        },
        3 => {
            4
        }
    };
}"#,
        );
    }

    #[test]
    fn test_pull_assignment_up_not_last_not_applicable() {
        check_assist_not_applicable(
            pull_assignment_up,
            r#"
fn foo() {
    let mut a = 1;

    if true {
        $0a = 2;
        b = a;
    } else {
        a = 3;
    }
}"#,
        )
    }

    #[test]
    fn test_pull_assignment_up_chained_if() {
        cov_mark::check!(test_pull_assignment_up_chained_if);
        check_assist(
            pull_assignment_up,
            r#"
fn foo() {
    let mut a = 1;

    if true {
        $0a = 2;
    } else if false {
        a = 3;
    } else {
        a = 4;
    }
}"#,
            r#"
fn foo() {
    let mut a = 1;

    a = if true {
        2
    } else if false {
        3
    } else {
        4
    };
}"#,
        );
    }

    #[test]
    fn test_pull_assignment_up_retains_stmts() {
        check_assist(
            pull_assignment_up,
            r#"
fn foo() {
    let mut a = 1;

    if true {
        let b = 2;
        $0a = 2;
    } else {
        let b = 3;
        a = 3;
    }
}"#,
            r#"
fn foo() {
    let mut a = 1;

    a = if true {
        let b = 2;
        2
    } else {
        let b = 3;
        3
    };
}"#,
        )
    }

    #[test]
    fn pull_assignment_up_let_stmt_not_applicable() {
        check_assist_not_applicable(
            pull_assignment_up,
            r#"
fn foo() {
    let mut a = 1;

    let b = if true {
        $0a = 2
    } else {
        a = 3
    };
}"#,
        )
    }

    #[test]
    fn pull_assignment_up_if_missing_assigment_not_applicable() {
        check_assist_not_applicable(
            pull_assignment_up,
            r#"
fn foo() {
    let mut a = 1;

    if true {
        $0a = 2;
    } else {}
}"#,
        )
    }

    #[test]
    fn pull_assignment_up_match_missing_assigment_not_applicable() {
        check_assist_not_applicable(
            pull_assignment_up,
            r#"
fn foo() {
    let mut a = 1;

    match 1 {
        1 => {
            $0a = 2;
        },
        2 => {
            a = 3;
        },
        3 => {},
    }
}"#,
        )
    }

    #[test]
    fn test_pull_assignment_up_field_assignment() {
        cov_mark::check!(test_pull_assignment_up_field_assignment);
        check_assist(
            pull_assignment_up,
            r#"
struct A(usize);

fn foo() {
    let mut a = A(1);

    if true {
        $0a.0 = 2;
    } else {
        a.0 = 3;
    }
}"#,
            r#"
struct A(usize);

fn foo() {
    let mut a = A(1);

    a.0 = if true {
        2
    } else {
        3
    };
}"#,
        )
    }

    #[test]
    fn test_pull_assignment_up_deref() {
        cov_mark::check!(test_pull_assignment_up_deref);
        check_assist(
            pull_assignment_up,
            r#"
fn foo() {
    let mut a = 1;
    let b = &mut a;

    if true {
        $0*b = 2;
    } else {
        *b = 3;
    }
}
"#,
            r#"
fn foo() {
    let mut a = 1;
    let b = &mut a;

    *b = if true {
        2
    } else {
        3
    };
}
"#,
        )
    }

    #[test]
    fn test_cant_pull_non_assignments() {
        cov_mark::check!(test_cant_pull_non_assignments);
        check_assist_not_applicable(
            pull_assignment_up,
            r#"
fn foo() {
    let mut a = 1;
    let b = &mut a;

    if true {
        $0*b + 2;
    } else {
        *b + 3;
    }
}
"#,
        )
    }
}
<|endoftext|>


<commit_before>//! This module contains an abstraction for HTTP clients as well as friendly-named re-exports of
//! client types that implement this trait.

use std::{future::Future, pin::Pin};

use async_trait::async_trait;
use bytes::BufMut;
use ruma_common::{
    api::{MatrixVersion, OutgoingRequest, SendAccessToken},
    UserId,
};

use crate::{add_user_id_to_query, ResponseError, ResponseResult};

#[cfg(feature = "hyper")]
mod hyper;
#[cfg(feature = "isahc")]
mod isahc;
#[cfg(feature = "reqwest")]
mod reqwest;

#[cfg(feature = "hyper")]
pub use self::hyper::Hyper;
#[cfg(feature = "hyper-native-tls")]
pub use self::hyper::HyperNativeTls;
#[cfg(feature = "hyper-rustls")]
pub use self::hyper::HyperRustls;
#[cfg(feature = "isahc")]
pub use self::isahc::Isahc;
#[cfg(feature = "reqwest")]
pub use self::reqwest::Reqwest;

/// An HTTP client that can be used to send requests to a Matrix homeserver.
#[async_trait]
pub trait HttpClient: Sync {
    /// The type to use for `try_into_http_request`.
    type RequestBody: Default + BufMut + Send;

    /// The type to use for `try_from_http_response`.
    type ResponseBody: AsRef<[u8]>;

    /// The error type for the `send_request` function.
    type Error: Send + Unpin;

    /// Send an `http::Request` to get back an `http::Response`.
    async fn send_http_request(
        &self,
        req: http::Request<Self::RequestBody>,
    ) -> Result<http::Response<Self::ResponseBody>, Self::Error>;
}

/// An HTTP client that has a default configuration.
pub trait DefaultConstructibleHttpClient: HttpClient {
    /// Creates a new HTTP client with default configuration.
    fn default() -> Self;
}

/// Convenience functionality on top of `HttpClient`.
///
/// If you want to build your own matrix client type instead of using `ruma_client::Client`, this
/// trait should make that relatively easy.
pub trait HttpClientExt: HttpClient {
    /// Send a strongly-typed matrix request to get back a strongly-typed response.
    // TODO: `R: 'a` bound should not be needed
    fn send_matrix_request<'a, R: OutgoingRequest + 'a>(
        &'a self,
        homeserver_url: &str,
        access_token: SendAccessToken<'_>,
        for_versions: &[MatrixVersion],
        request: R,
    ) -> Pin<Box<dyn Future<Output = ResponseResult<Self, R>> + 'a>> {
        self.send_customized_matrix_request(
            homeserver_url,
            access_token,
            for_versions,
            request,
            |_| Ok(()),
        )
    }

    /// Turn a strongly-typed matrix request into an `http::Request`, customize it and send it to
    /// get back a strongly-typed response.
    // TODO: `R: 'a` and `F: 'a` should not be needed
    fn send_customized_matrix_request<'a, R, F>(
        &'a self,
        homeserver_url: &str,
        access_token: SendAccessToken<'_>,
        for_versions: &[MatrixVersion],
        request: R,
        customize: F,
    ) -> Pin<Box<dyn Future<Output = ResponseResult<Self, R>> + 'a>>
    where
        R: OutgoingRequest + 'a,
        F: FnOnce(&mut http::Request<Self::RequestBody>) -> Result<(), ResponseError<Self, R>> + 'a,
    {
        Box::pin(crate::send_customized_request(
            self,
            homeserver_url,
            access_token,
            for_versions,
            request,
            customize,
        ))
    }

    /// Turn a strongly-typed matrix request into an `http::Request`, add a `user_id` query
    /// parameter to it and send it to get back a strongly-typed response.
    ///
    /// This method is meant to be used by application services when interacting with the
    /// client-server API.
    fn send_matrix_request_as<'a, R: OutgoingRequest + 'a>(
        &'a self,
        homeserver_url: &str,
        access_token: SendAccessToken<'_>,
        for_versions: &[MatrixVersion],
        user_id: &'a UserId,
        request: R,
    ) -> Pin<Box<dyn Future<Output = ResponseResult<Self, R>> + 'a>> {
        self.send_customized_matrix_request(
            homeserver_url,
            access_token,
            for_versions,
            request,
            add_user_id_to_query::<Self, R>(user_id),
        )
    }
}

#[async_trait]
impl<T: HttpClient> HttpClientExt for T {}

#[doc(hidden)]
#[derive(Debug)]
#[allow(clippy::exhaustive_structs)]
pub struct Dummy;

#[async_trait]
impl HttpClient for Dummy {
    type RequestBody = Vec<u8>;
    type ResponseBody = Vec<u8>;
    type Error = ();

    async fn send_http_request(
        &self,
        _req: http::Request<Self::RequestBody>,
    ) -> Result<http::Response<Self::ResponseBody>, Self::Error> {
        unimplemented!("this client only exists to allow doctests to compile")
    }
}

impl DefaultConstructibleHttpClient for Dummy {
    fn default() -> Self {
        Dummy
    }
}
<commit_msg>client: Mark futures returned by send_matrix_request as Send<commit_after>//! This module contains an abstraction for HTTP clients as well as friendly-named re-exports of
//! client types that implement this trait.

use std::{future::Future, pin::Pin};

use async_trait::async_trait;
use bytes::BufMut;
use ruma_common::{
    api::{MatrixVersion, OutgoingRequest, SendAccessToken},
    UserId,
};

use crate::{add_user_id_to_query, ResponseError, ResponseResult};

#[cfg(feature = "hyper")]
mod hyper;
#[cfg(feature = "isahc")]
mod isahc;
#[cfg(feature = "reqwest")]
mod reqwest;

#[cfg(feature = "hyper")]
pub use self::hyper::Hyper;
#[cfg(feature = "hyper-native-tls")]
pub use self::hyper::HyperNativeTls;
#[cfg(feature = "hyper-rustls")]
pub use self::hyper::HyperRustls;
#[cfg(feature = "isahc")]
pub use self::isahc::Isahc;
#[cfg(feature = "reqwest")]
pub use self::reqwest::Reqwest;

/// An HTTP client that can be used to send requests to a Matrix homeserver.
#[async_trait]
pub trait HttpClient: Sync {
    /// The type to use for `try_into_http_request`.
    type RequestBody: Default + BufMut + Send;

    /// The type to use for `try_from_http_response`.
    type ResponseBody: AsRef<[u8]>;

    /// The error type for the `send_request` function.
    type Error: Send + Unpin;

    /// Send an `http::Request` to get back an `http::Response`.
    async fn send_http_request(
        &self,
        req: http::Request<Self::RequestBody>,
    ) -> Result<http::Response<Self::ResponseBody>, Self::Error>;
}

/// An HTTP client that has a default configuration.
pub trait DefaultConstructibleHttpClient: HttpClient {
    /// Creates a new HTTP client with default configuration.
    fn default() -> Self;
}

/// Convenience functionality on top of `HttpClient`.
///
/// If you want to build your own matrix client type instead of using `ruma_client::Client`, this
/// trait should make that relatively easy.
pub trait HttpClientExt: HttpClient {
    /// Send a strongly-typed matrix request to get back a strongly-typed response.
    // TODO: `R: 'a` bound should not be needed
    fn send_matrix_request<'a, R: OutgoingRequest + 'a>(
        &'a self,
        homeserver_url: &str,
        access_token: SendAccessToken<'_>,
        for_versions: &[MatrixVersion],
        request: R,
    ) -> Pin<Box<dyn Future<Output = ResponseResult<Self, R>> + 'a + Send>> {
        self.send_customized_matrix_request(
            homeserver_url,
            access_token,
            for_versions,
            request,
            |_| Ok(()),
        )
    }

    /// Turn a strongly-typed matrix request into an `http::Request`, customize it and send it to
    /// get back a strongly-typed response.
    // TODO: `R: 'a` and `F: 'a` should not be needed
    fn send_customized_matrix_request<'a, R, F>(
        &'a self,
        homeserver_url: &str,
        access_token: SendAccessToken<'_>,
        for_versions: &[MatrixVersion],
        request: R,
        customize: F,
    ) -> Pin<Box<dyn Future<Output = ResponseResult<Self, R>> + 'a + Send>>
    where
        R: OutgoingRequest + 'a,
        F: FnOnce(&mut http::Request<Self::RequestBody>) -> Result<(), ResponseError<Self, R>> + 'a,
    {
        Box::pin(crate::send_customized_request(
            self,
            homeserver_url,
            access_token,
            for_versions,
            request,
            customize,
        ))
    }

    /// Turn a strongly-typed matrix request into an `http::Request`, add a `user_id` query
    /// parameter to it and send it to get back a strongly-typed response.
    ///
    /// This method is meant to be used by application services when interacting with the
    /// client-server API.
    fn send_matrix_request_as<'a, R: OutgoingRequest + 'a>(
        &'a self,
        homeserver_url: &str,
        access_token: SendAccessToken<'_>,
        for_versions: &[MatrixVersion],
        user_id: &'a UserId,
        request: R,
    ) -> Pin<Box<dyn Future<Output = ResponseResult<Self, R>> + 'a>> {
        self.send_customized_matrix_request(
            homeserver_url,
            access_token,
            for_versions,
            request,
            add_user_id_to_query::<Self, R>(user_id),
        )
    }
}

#[async_trait]
impl<T: HttpClient> HttpClientExt for T {}

#[doc(hidden)]
#[derive(Debug)]
#[allow(clippy::exhaustive_structs)]
pub struct Dummy;

#[async_trait]
impl HttpClient for Dummy {
    type RequestBody = Vec<u8>;
    type ResponseBody = Vec<u8>;
    type Error = ();

    async fn send_http_request(
        &self,
        _req: http::Request<Self::RequestBody>,
    ) -> Result<http::Response<Self::ResponseBody>, Self::Error> {
        unimplemented!("this client only exists to allow doctests to compile")
    }
}

impl DefaultConstructibleHttpClient for Dummy {
    fn default() -> Self {
        Dummy
    }
}
<|endoftext|>


<commit_before>/// Internal namespace.
pub mod internal
{

  ///
  /// Internals of a test suite.
  ///

  #[ macro_export ]
  macro_rules! impls
  {

    () => {};

    (
      $( #[ $Meta : meta ] )*
      $Vis : vis
      fn $Name : ident
      $( $Rest : tt )*
    )
    =>
    {
      $crate::impls!
      {
        @DEFINE_FN
        @META{ $( #[ $Meta ] )* }
        @VIS{ $Vis }
        @NAME{ $Name }
        @INPUT{ () }
        @OUTPUT{}
        @BLOCK{ {} }
        @REST
          $( #[ $Meta ] )*
          $Vis fn $Name
          $( $Rest )*
      }
    };

    (
      @DEFINE_FN
      @META{ $( #[ $Meta : meta ] )* }
      @VIS{ $Vis : vis }
      @NAME{ $Name : ident }
      @INPUT{ $Input : tt }
      @OUTPUT{ $( -> $Output : ty )? }
      @BLOCK{ $Block : block }
      @REST
        $Item : item
        $( $Rest : tt )*
    )
    =>
    {
      #[ deny( unused_macros ) ]
      macro_rules! $Name
      {
        () =>
        {
          $Item
        };
        // ( @AS $Name : ident ) =>
        // {
        //   $( #[ $Meta ] )*
        //   fn $Name
        // };
      }

      $crate::impls!
      {
        $( $Rest )*
      }
    };

  }

  ///
  /// Index of items.
  ///

  #[ macro_export ]
  macro_rules! index
  {

    () => { };

    (
      $Name : ident ,
      $( $Rest : tt )*
    )
    =>
    {
      $Name!();
      $crate::index!( $( $Rest )* );
    };

  }

  ///
  /// Index of items.
  ///

  #[ macro_export ]
  macro_rules! ignore_macro
  {

    () => { };

    (
      $Name : ident ,
      $( $Rest : tt )*
    )
    =>
    {
      $Name!();
      stringify!( $crate::index!( $( $Rest )* ) );
    };

  }

  pub use impls;
  pub use index;
  pub use ignore_macro;
}

/// Exposed namespace of the module.
pub mod exposed
{
  // use super::internal as i;
  pub use super::prelude::*;
}

pub use exposed::*;

/// Prelude to use: `use wtools::prelude::*`.
pub mod prelude
{
  use super::internal as i;
  pub use i::impls;
  pub use i::index;
  pub use i::ignore_macro;
}

//   #[ macro_export ]
//   macro_rules! impls
//   {
//
//     () => {};
//
//     (
//       $( #[ $Meta : meta ] )*
//       $Vis : vis
//       fn $Name : ident
//       $( $Rest : tt )*
//     )
//     =>
//     {
//       $crate::impls!
//       {
//         @DEFINE_FN
//         @META{ $( #[ $Meta ] )* }
//         @VIS{ $Vis }
//         @NAME{ $Name }
//         @INPUT{ () }
//         @OUTPUT{}
//         @BLOCK{ {} }
//         @REST
//           $( #[ $Meta ] )*
//           $Vis fn $Name
//           $( $Rest )*
//       }
//     };
//
//     (
//       @DEFINE_FN
//       @META{ $( #[ $Meta : meta ] )* }
//       @VIS{ $Vis : vis }
//       @NAME{ $Name : ident }
//       @INPUT{ $Input : tt }
//       @OUTPUT{ $( -> $Output : ty )? }
//       @BLOCK{ $Block : block }
//       @REST
//         $Item : item
//         $( $Rest : tt )*
//     )
//     =>
//     {
//       #[ deny( unused_macros ) ]
//       macro_rules! $Name
//       {
//         () =>
//         {
//           $Item
//         };
//         // ( @AS $Name : ident ) =>
//         // {
//         //   $( #[ $Meta ] )*
//         //   fn $Name
//         // };
//       }
//
//       $crate::impls!
//       {
//         $( $Rest )*
//       }
//     };
//
//   }

//   #[ macro_export ]
//   macro_rules! impls
//   {
//
//     () => {};
//
//     (
//       $( #[ $Meta : meta ] )*
//       $Vis : vis
//       fn $Name : ident
//       $( $Rest : tt )*
//     )
//     =>
//     {
//       $crate::impls!
//       {
//         @DEFINE_FN
//         @META $( #[ $Meta ] )*
//         @VIS{ $Vis }
//         @NAME $Name
//         // @INPUT ()
//         // @OUTPUT
//         // @BLOCK {}
//         @REST
//           $( #[ $Meta ] )*
//           $Vis fn $Name
//           $( $Rest )*
//       }
//     };
//
//     // (
//     //   $( #[ $Meta : meta ] )*
//     //   fn $Name : ident
//     //   // $( < $( $ParamName : ident : $ParamType : ty ),* > )?
//     //   // $Input : tt
//     //   // $( -> $Output : ty )?
//     //   // $Block : block
//     //   $( $Rest : tt )*
//     // )
//     // =>
//     // {
//     //   $crate::impls!
//     //   {
//     //     @DEFINE_FN
//     //     @META $( #[ $Meta ] )*
//     //     @VIS
//     //     @NAME $Name
//     //     @INPUT ()
//     //     @OUTPUT
//     //     @BLOCK {}
//     //     @REST
//     //       $( #[ $Meta ] )*
//     //       fn $Name
//     //       // $Input $( -> $Output )?
//     //       // $Block
//     //       $( $Rest )*
//     //   }
//     // };
//
// //     (
// //       $( #[ $Meta : meta ] )*
// //       pub
// //       fn $Name : ident
// //       $Input : tt
// //       $( -> $Output : ty )?
// //       $Block : block
// //       $( $Rest : tt )*
// //     )
// //     =>
// //     {
// //       $crate::impls!
// //       {
// //         @DEFINE_FN
// //         @META $( #[ $Meta ] )*
// //         @VIS pub
// //         @NAME $Name
// //         @INPUT $Input
// //         @OUTPUT $( -> $Output )?
// //         @BLOCK $Block
// //         @REST
// //           $( #[ $Meta ] )*
// //           pub fn $Name $Input $( -> $Output )?
// //           $Block
// //           $( $Rest )*
// //       }
// //     };
// //
// //     (
// //       $( #[ $Meta : meta ] )*
// //       fn $Name : ident
// //       $Input : tt
// //       $( -> $Output : ty )?
// //       $Block : block
// //       $( $Rest : tt )*
// //     )
// //     =>
// //     {
// //       $crate::impls!
// //       {
// //         @DEFINE_FN
// //         @META $( #[ $Meta ] )*
// //         @VIS
// //         @NAME $Name
// //         @INPUT $Input
// //         @OUTPUT $( -> $Output )?
// //         @BLOCK $Block
// //         @REST
// //           $( #[ $Meta ] )*
// //           fn $Name $Input $( -> $Output )?
// //           $Block
// //           $( $Rest )*
// //       }
// //     };
//
//     (
//       @DEFINE_FN
//       @META $( #[ $Meta : meta ] )*
//       @VIS{ $Vis : vis }
//       @NAME $Name : ident
//       // @INPUT $Input : tt
//       // @OUTPUT $( -> $Output : ty )?
//       // @BLOCK $Block : block
//       @REST
//         $Item : item
//         $( $Rest : tt )*
//     )
//     =>
//     {
//       // #[ deny( unused_macros ) ]
//       macro_rules! $Name
//       {
//         () =>
//         {
//           $Item
//         };
//         // ( @AS $Name : ident ) =>
//         // {
//         //   $( #[ $Meta ] )*
//         //   fn $Name
//         // };
//       }
//
//       $crate::impls!
//       {
//         $( $Rest )*
//       }
//     };
//
//   }
<commit_msg>extend macro impls<commit_after>/// Internal namespace.
pub mod internal
{

  ///
  /// Internals of a test suite.
  ///

  #[ macro_export ]
  macro_rules! impls
  {

    () => {};

    (
      $( #[ $Meta : meta ] )*
      $Vis : vis
      fn $Name : ident
      $( $Rest : tt )*
    )
    =>
    {
      $crate::impls!
      {
        @DEFINE_FN
        @META{ $( #[ $Meta ] )* }
        @VIS{ $Vis }
        @NAME{ $Name }
        @INPUT{ () }
        @OUTPUT{}
        @BLOCK{ {} }
        @REST
          $( #[ $Meta ] )*
          $Vis fn $Name
          $( $Rest )*
      }
    };

    (
      @DEFINE_FN
      @META{ $( #[ $Meta : meta ] )* }
      @VIS{ $Vis : vis }
      @NAME{ $Name : ident }
      @INPUT{ $Input : tt }
      @OUTPUT{ $( -> $Output : ty )? }
      @BLOCK{ $Block : block }
      @REST
        $Item : item
        $( $Rest : tt )*
    )
    =>
    {
      // #[ deny( unused_macros ) ]
      macro_rules! $Name
      {
        () =>
        {
          $Item
        };
        // ( @AS $Name : ident ) =>
        // {
        //   $( #[ $Meta ] )*
        //   fn $Name
        // };
      }

      $crate::impls!
      {
        $( $Rest )*
      }
    };

  }

  ///
  /// Index of items.
  ///

  #[ macro_export ]
  macro_rules! index
  {

    () => { };

    (
      $Name : ident ,
      $( $Rest : tt )*
    )
    =>
    {
      $Name!();
      $crate::index!( $( $Rest )* );
    };

  }

  ///
  /// Index of items.
  ///

  #[ macro_export ]
  macro_rules! ignore_macro
  {

    () => { };

    (
      $Name : ident ,
      $( $Rest : tt )*
    )
    =>
    {
      $Name!();
      stringify!( $crate::index!( $( $Rest )* ) );
    };

  }

  pub use impls;
  pub use index;
  pub use ignore_macro;
}

/// Exposed namespace of the module.
pub mod exposed
{
  // use super::internal as i;
  pub use super::prelude::*;
}

pub use exposed::*;

/// Prelude to use: `use wtools::prelude::*`.
pub mod prelude
{
  use super::internal as i;
  pub use i::impls;
  pub use i::index;
  pub use i::ignore_macro;
}

//   #[ macro_export ]
//   macro_rules! impls
//   {
//
//     () => {};
//
//     (
//       $( #[ $Meta : meta ] )*
//       $Vis : vis
//       fn $Name : ident
//       $( $Rest : tt )*
//     )
//     =>
//     {
//       $crate::impls!
//       {
//         @DEFINE_FN
//         @META{ $( #[ $Meta ] )* }
//         @VIS{ $Vis }
//         @NAME{ $Name }
//         @INPUT{ () }
//         @OUTPUT{}
//         @BLOCK{ {} }
//         @REST
//           $( #[ $Meta ] )*
//           $Vis fn $Name
//           $( $Rest )*
//       }
//     };
//
//     (
//       @DEFINE_FN
//       @META{ $( #[ $Meta : meta ] )* }
//       @VIS{ $Vis : vis }
//       @NAME{ $Name : ident }
//       @INPUT{ $Input : tt }
//       @OUTPUT{ $( -> $Output : ty )? }
//       @BLOCK{ $Block : block }
//       @REST
//         $Item : item
//         $( $Rest : tt )*
//     )
//     =>
//     {
//       #[ deny( unused_macros ) ]
//       macro_rules! $Name
//       {
//         () =>
//         {
//           $Item
//         };
//         // ( @AS $Name : ident ) =>
//         // {
//         //   $( #[ $Meta ] )*
//         //   fn $Name
//         // };
//       }
//
//       $crate::impls!
//       {
//         $( $Rest )*
//       }
//     };
//
//   }

//   #[ macro_export ]
//   macro_rules! impls
//   {
//
//     () => {};
//
//     (
//       $( #[ $Meta : meta ] )*
//       $Vis : vis
//       fn $Name : ident
//       $( $Rest : tt )*
//     )
//     =>
//     {
//       $crate::impls!
//       {
//         @DEFINE_FN
//         @META $( #[ $Meta ] )*
//         @VIS{ $Vis }
//         @NAME $Name
//         // @INPUT ()
//         // @OUTPUT
//         // @BLOCK {}
//         @REST
//           $( #[ $Meta ] )*
//           $Vis fn $Name
//           $( $Rest )*
//       }
//     };
//
//     // (
//     //   $( #[ $Meta : meta ] )*
//     //   fn $Name : ident
//     //   // $( < $( $ParamName : ident : $ParamType : ty ),* > )?
//     //   // $Input : tt
//     //   // $( -> $Output : ty )?
//     //   // $Block : block
//     //   $( $Rest : tt )*
//     // )
//     // =>
//     // {
//     //   $crate::impls!
//     //   {
//     //     @DEFINE_FN
//     //     @META $( #[ $Meta ] )*
//     //     @VIS
//     //     @NAME $Name
//     //     @INPUT ()
//     //     @OUTPUT
//     //     @BLOCK {}
//     //     @REST
//     //       $( #[ $Meta ] )*
//     //       fn $Name
//     //       // $Input $( -> $Output )?
//     //       // $Block
//     //       $( $Rest )*
//     //   }
//     // };
//
// //     (
// //       $( #[ $Meta : meta ] )*
// //       pub
// //       fn $Name : ident
// //       $Input : tt
// //       $( -> $Output : ty )?
// //       $Block : block
// //       $( $Rest : tt )*
// //     )
// //     =>
// //     {
// //       $crate::impls!
// //       {
// //         @DEFINE_FN
// //         @META $( #[ $Meta ] )*
// //         @VIS pub
// //         @NAME $Name
// //         @INPUT $Input
// //         @OUTPUT $( -> $Output )?
// //         @BLOCK $Block
// //         @REST
// //           $( #[ $Meta ] )*
// //           pub fn $Name $Input $( -> $Output )?
// //           $Block
// //           $( $Rest )*
// //       }
// //     };
// //
// //     (
// //       $( #[ $Meta : meta ] )*
// //       fn $Name : ident
// //       $Input : tt
// //       $( -> $Output : ty )?
// //       $Block : block
// //       $( $Rest : tt )*
// //     )
// //     =>
// //     {
// //       $crate::impls!
// //       {
// //         @DEFINE_FN
// //         @META $( #[ $Meta ] )*
// //         @VIS
// //         @NAME $Name
// //         @INPUT $Input
// //         @OUTPUT $( -> $Output )?
// //         @BLOCK $Block
// //         @REST
// //           $( #[ $Meta ] )*
// //           fn $Name $Input $( -> $Output )?
// //           $Block
// //           $( $Rest )*
// //       }
// //     };
//
//     (
//       @DEFINE_FN
//       @META $( #[ $Meta : meta ] )*
//       @VIS{ $Vis : vis }
//       @NAME $Name : ident
//       // @INPUT $Input : tt
//       // @OUTPUT $( -> $Output : ty )?
//       // @BLOCK $Block : block
//       @REST
//         $Item : item
//         $( $Rest : tt )*
//     )
//     =>
//     {
//       // #[ deny( unused_macros ) ]
//       macro_rules! $Name
//       {
//         () =>
//         {
//           $Item
//         };
//         // ( @AS $Name : ident ) =>
//         // {
//         //   $( #[ $Meta ] )*
//         //   fn $Name
//         // };
//       }
//
//       $crate::impls!
//       {
//         $( $Rest )*
//       }
//     };
//
//   }
<|endoftext|>


<commit_before>//! This module contains the high-level DICOM abstraction trait.
//! At this level, objects are comparable to a lazy dictionary of elements,
//! in which some of them can be DICOM objects themselves.
//! The end user should prefer using this abstraction when dealing with DICOM objects.
//!
//! # Examples
//!
//! Fetching an element by tag:
//!
//! ```
//! # use dicom_core::object::DicomObject;
//! # use dicom_core::error::Result;
//! use dicom_core::data::Tag;
//! # fn something<T: DicomObject>(obj: T) -> Result<()> {
//! let e = obj.element(Tag(0x0002, 0x0002))?;
//! # Ok(())
//! # }
//! ```
//!
use data::Header;
use data::Tag;
use error::Result;

pub mod lazy;
pub mod mem;
pub mod pixeldata;

/// Trait type for a DICOM object.
/// This is a high-level abstraction where an object is accessed and
/// manipulated as dictionary of entries indexed by tags, which in
/// turn may contain a DICOM object.
///
/// This trait interface is experimental and prone to sudden changes.
pub trait DicomObject {
    type Element: Header; // TODO change constraint

    /// Retrieve a particular DICOM element by its tag.
    fn element(&self, tag: Tag) -> Result<Self::Element>;

    /// Retrieve a particular DICOM element by its name.
    fn element_by_name(&self, name: &str) -> Result<Self::Element>;

    // TODO moar
}
<commit_msg>Introduce Root DICOM object<commit_after>//! This module contains the high-level DICOM abstraction trait.
//! At this level, objects are comparable to a lazy dictionary of elements,
//! in which some of them can be DICOM objects themselves.
//! The end user should prefer using this abstraction when dealing with DICOM objects.
//!
//! # Examples
//!
//! Fetching an element by tag:
//!
//! ```
//! # use dicom_core::object::DicomObject;
//! # use dicom_core::error::Result;
//! use dicom_core::data::Tag;
//! # fn something<T: DicomObject>(obj: T) -> Result<()> {
//! let e = obj.element(Tag(0x0002, 0x0002))?;
//! # Ok(())
//! # }
//! ```
//!
use data::Header;
use data::Tag;
use error::Result;
use meta::DicomMetaTable;

pub mod lazy;
pub mod mem;
pub mod pixeldata;

/// Trait type for a DICOM object.
/// This is a high-level abstraction where an object is accessed and
/// manipulated as dictionary of entries indexed by tags, which in
/// turn may contain a DICOM object.
///
/// This trait interface is experimental and prone to sudden changes.
pub trait DicomObject {
    type Element: Header; // TODO change constraint

    /// Retrieve a particular DICOM element by its tag.
    fn element(&self, tag: Tag) -> Result<Self::Element>;

    /// Retrieve a particular DICOM element by its name.
    fn element_by_name(&self, name: &str) -> Result<Self::Element>;

    // TODO moar
}

/** A root DICOM object contains additional meta information about the object
 * (such as the DICOM file's meta header).
 */
#[derive(Debug, Clone, PartialEq)]
pub struct RootDicomObject<T> {
    meta: DicomMetaTable,
    obj: T,
}

impl<T> RootDicomObject<T> {
    /// Retrieve the processed meta header table.
    pub fn meta(&self) -> &DicomMetaTable {
        &self.meta
    }
}

impl<T> ::std::ops::Deref for RootDicomObject<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.obj
    }
}

impl<T> ::std::ops::DerefMut for RootDicomObject<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.obj
    }
}

impl<T> DicomObject for RootDicomObject<T>
where
    T: DicomObject
{
    type Element = <T as DicomObject>::Element;

    fn element(&self, tag: Tag) -> Result<Self::Element> {
        self.obj.element(tag)
    }

    fn element_by_name(&self, name: &str) -> Result<Self::Element> {
        self.obj.element_by_name(name)
    }
}

impl<'a, T: 'a> DicomObject for &'a RootDicomObject<T>
where
    T: DicomObject
{
    type Element = <T as DicomObject>::Element;

    fn element(&self, tag: Tag) -> Result<Self::Element> {
        self.obj.element(tag)
    }

    fn element_by_name(&self, name: &str) -> Result<Self::Element> {
        self.obj.element_by_name(name)
    }
}
<|endoftext|>


<commit_before>//! Unix terminal formatter and drain for slog-rs
//!
//! ```
//! #[macro_use]
//! extern crate slog;
//! extern crate slog_term;
//!
//! use slog::*;
//! use slog::drain::IntoLogger;
//!
//! fn main() {
//!     let root = slog_term::async_stderr().into_logger(o!("build-id" => "8dfljdf"));
//! }
//! ```
#![warn(missing_docs)]

#[macro_use]
extern crate slog;
extern crate isatty;
extern crate ansi_term;
extern crate rustc_serialize;

use std::io;
use rustc_serialize::hex::ToHex;

use ansi_term::Colour;
use isatty::{stderr_isatty, stdout_isatty};

use slog::logger::RecordInfo;
use slog::drain::{Streamer, AsyncStreamer};
use slog::{Level, OwnedKeyValueNode};
use slog::format::Format as SlogFormat;

/// Format formatting with optional color support
pub struct Format {
    color: bool,
}

impl Format {
    /// New Format format that prints using color
    pub fn colored() -> Self {
        Format { color: true }
    }

    /// New Format format that prints without using color
    pub fn plain() -> Self {
        Format { color: false }
    }
}

fn severity_to_color(lvl: Level) -> u8 {
    match lvl {
        Level::Critical => 1,
        Level::Error => 9,
        Level::Warning => 3,
        Level::Info => 2,
        Level::Debug => 10,
        Level::Trace => 11,
    }
}

struct Serializer<W> {
    io: W,
    key: ansi_term::Style,
}

impl<W: io::Write> Serializer<W> {
    fn new(io: W, color: bool) -> Self {
        Serializer {
            io: io,
            key: if color {
                ansi_term::Style::new().bold()
            } else {
                ansi_term::Style::new()
            },
        }
    }

    fn maybe_comma(&mut self) -> std::io::Result<()> {
        // since we don't use it for first element, add comma
        // unconditionally
        try!(write!(self.io, ", "));
        Ok(())
    }

    fn into_inner(self) -> W {
        self.io
    }
}

impl<W: io::Write> slog::ser::Serializer for Serializer<W> {
    fn emit_none(&mut self, key: &str) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), "None"));
        Ok(())
    }
    fn emit_unit(&mut self, key: &str) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: ()", self.key.paint(key)));
        Ok(())
    }

    fn emit_bool(&mut self, key: &str, val: bool) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }

    fn emit_char(&mut self, key: &str, val: char) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }
    fn emit_bytes(&mut self, key: &str, val: &[u8]) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val.to_hex()));
        Ok(())
    }

    fn emit_usize(&mut self, key: &str, val: usize) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }
    fn emit_isize(&mut self, key: &str, val: isize) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }

    fn emit_u8(&mut self, key: &str, val: u8) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }
    fn emit_i8(&mut self, key: &str, val: i8) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }
    fn emit_u16(&mut self, key: &str, val: u16) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }
    fn emit_i16(&mut self, key: &str, val: i16) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }
    fn emit_u32(&mut self, key: &str, val: u32) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }
    fn emit_i32(&mut self, key: &str, val: i32) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }
    fn emit_f32(&mut self, key: &str, val: f32) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }
    fn emit_u64(&mut self, key: &str, val: u64) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }
    fn emit_i64(&mut self, key: &str, val: i64) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }
    fn emit_f64(&mut self, key: &str, val: f64) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }
    fn emit_str(&mut self, key: &str, val: &str) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }
}

impl SlogFormat for Format {
    fn format(&self,
              io: &mut io::Write,
              info: &RecordInfo,
              logger_values: &OwnedKeyValueNode)
              -> slog::format::Result<()> {
        let level_color = Colour::Fixed(severity_to_color(info.level()));
        let bold = ansi_term::Style::new().bold();

        if self.color {
            let _ = try!(write!(io,
                                "{} {} {}",
                                info.ts().format("%b %d %H:%M:%S%.3f"),
                                level_color.paint(info.level().as_short_str()),
                                bold.paint(info.msg()).to_string()));
        } else {
            let _ = try!(write!(io,
                                "{} {} {}",
                                info.ts().format("%b %d %H:%M:%S%.3f"),
                                info.level().as_short_str(),
                                info.msg()));
        }

        let mut serializer = Serializer::new(io, self.color);

        for &(ref k, ref v) in logger_values.iter() {
            try!(v.serialize(info, k, &mut serializer));
        }

        for &(k, v) in info.values() {
            try!(v.serialize(info, k, &mut serializer));
        }
        let mut io = serializer.into_inner();

        let _ = try!(write!(io, "\n"));

        Ok(())
    }
}

/// Short for `Format::colored()`
pub fn format_colored() -> Format {
    Format::colored()
}

/// Short for `Format::plain()`
pub fn format_plain() -> Format {
    Format::plain()
}

/// Drain to `stdout`
///
/// Automatically using color if printing to tty
pub fn stdout() -> Streamer<io::Stdout, Format> {
    Streamer::new(io::stdout(),
                  if stdout_isatty() {
                      Format::colored()
                  } else {
                      Format::plain()
                  })
}

/// Drain to `stderr`
///
/// Automatically using color if output goes to tty
pub fn stderr() -> Streamer<io::Stderr, Format> {
    Streamer::new(io::stderr(),
                  if stderr_isatty() {
                      Format::colored()
                  } else {
                      Format::plain()
                  })
}

/// Asynchronous drain to `stdout`
///
/// Automatically using color if printing to tty
pub fn async_stdout() -> AsyncStreamer<Format> {
    AsyncStreamer::new(io::stdout(),
                       if stdout_isatty() {
                           Format::colored()
                       } else {
                           Format::plain()
                       })
}

/// Asynchronos drain to `stderr`
///
/// Automatically using color if output goes to tty
pub fn async_stderr() -> AsyncStreamer<Format> {
    AsyncStreamer::new(io::stderr(),
                       if stderr_isatty() {
                           Format::colored()
                       } else {
                           Format::plain()
                       })
}
<commit_msg>Tune `term`s debug and trace colors<commit_after>//! Unix terminal formatter and drain for slog-rs
//!
//! ```
//! #[macro_use]
//! extern crate slog;
//! extern crate slog_term;
//!
//! use slog::*;
//! use slog::drain::IntoLogger;
//!
//! fn main() {
//!     let root = slog_term::async_stderr().into_logger(o!("build-id" => "8dfljdf"));
//! }
//! ```
#![warn(missing_docs)]

#[macro_use]
extern crate slog;
extern crate isatty;
extern crate ansi_term;
extern crate rustc_serialize;

use std::io;
use rustc_serialize::hex::ToHex;

use ansi_term::Colour;
use isatty::{stderr_isatty, stdout_isatty};

use slog::logger::RecordInfo;
use slog::drain::{Streamer, AsyncStreamer};
use slog::{Level, OwnedKeyValueNode};
use slog::format::Format as SlogFormat;

/// Format formatting with optional color support
pub struct Format {
    color: bool,
}

impl Format {
    /// New Format format that prints using color
    pub fn colored() -> Self {
        Format { color: true }
    }

    /// New Format format that prints without using color
    pub fn plain() -> Self {
        Format { color: false }
    }
}

fn severity_to_color(lvl: Level) -> u8 {
    match lvl {
        Level::Critical => 1,
        Level::Error => 9,
        Level::Warning => 3,
        Level::Info => 2,
        Level::Debug => 6,
        Level::Trace => 4,
    }
}

struct Serializer<W> {
    io: W,
    key: ansi_term::Style,
}

impl<W: io::Write> Serializer<W> {
    fn new(io: W, color: bool) -> Self {
        Serializer {
            io: io,
            key: if color {
                ansi_term::Style::new().bold()
            } else {
                ansi_term::Style::new()
            },
        }
    }

    fn maybe_comma(&mut self) -> std::io::Result<()> {
        // since we don't use it for first element, add comma
        // unconditionally
        try!(write!(self.io, ", "));
        Ok(())
    }

    fn into_inner(self) -> W {
        self.io
    }
}

impl<W: io::Write> slog::ser::Serializer for Serializer<W> {
    fn emit_none(&mut self, key: &str) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), "None"));
        Ok(())
    }
    fn emit_unit(&mut self, key: &str) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: ()", self.key.paint(key)));
        Ok(())
    }

    fn emit_bool(&mut self, key: &str, val: bool) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }

    fn emit_char(&mut self, key: &str, val: char) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }
    fn emit_bytes(&mut self, key: &str, val: &[u8]) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val.to_hex()));
        Ok(())
    }

    fn emit_usize(&mut self, key: &str, val: usize) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }
    fn emit_isize(&mut self, key: &str, val: isize) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }

    fn emit_u8(&mut self, key: &str, val: u8) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }
    fn emit_i8(&mut self, key: &str, val: i8) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }
    fn emit_u16(&mut self, key: &str, val: u16) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }
    fn emit_i16(&mut self, key: &str, val: i16) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }
    fn emit_u32(&mut self, key: &str, val: u32) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }
    fn emit_i32(&mut self, key: &str, val: i32) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }
    fn emit_f32(&mut self, key: &str, val: f32) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }
    fn emit_u64(&mut self, key: &str, val: u64) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }
    fn emit_i64(&mut self, key: &str, val: i64) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }
    fn emit_f64(&mut self, key: &str, val: f64) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }
    fn emit_str(&mut self, key: &str, val: &str) -> slog::ser::Result<()> {
        try!(self.maybe_comma());
        try!(write!(self.io, "{}: {}", self.key.paint(key), val));
        Ok(())
    }
}

impl SlogFormat for Format {
    fn format(&self,
              io: &mut io::Write,
              info: &RecordInfo,
              logger_values: &OwnedKeyValueNode)
              -> slog::format::Result<()> {
        let level_color = Colour::Fixed(severity_to_color(info.level()));
        let bold = ansi_term::Style::new().bold();

        if self.color {
            let _ = try!(write!(io,
                                "{} {} {}",
                                info.ts().format("%b %d %H:%M:%S%.3f"),
                                level_color.paint(info.level().as_short_str()),
                                bold.paint(info.msg()).to_string()));
        } else {
            let _ = try!(write!(io,
                                "{} {} {}",
                                info.ts().format("%b %d %H:%M:%S%.3f"),
                                info.level().as_short_str(),
                                info.msg()));
        }

        let mut serializer = Serializer::new(io, self.color);

        for &(ref k, ref v) in logger_values.iter() {
            try!(v.serialize(info, k, &mut serializer));
        }

        for &(k, v) in info.values() {
            try!(v.serialize(info, k, &mut serializer));
        }
        let mut io = serializer.into_inner();

        let _ = try!(write!(io, "\n"));

        Ok(())
    }
}

/// Short for `Format::colored()`
pub fn format_colored() -> Format {
    Format::colored()
}

/// Short for `Format::plain()`
pub fn format_plain() -> Format {
    Format::plain()
}

/// Drain to `stdout`
///
/// Automatically using color if printing to tty
pub fn stdout() -> Streamer<io::Stdout, Format> {
    Streamer::new(io::stdout(),
                  if stdout_isatty() {
                      Format::colored()
                  } else {
                      Format::plain()
                  })
}

/// Drain to `stderr`
///
/// Automatically using color if output goes to tty
pub fn stderr() -> Streamer<io::Stderr, Format> {
    Streamer::new(io::stderr(),
                  if stderr_isatty() {
                      Format::colored()
                  } else {
                      Format::plain()
                  })
}

/// Asynchronous drain to `stdout`
///
/// Automatically using color if printing to tty
pub fn async_stdout() -> AsyncStreamer<Format> {
    AsyncStreamer::new(io::stdout(),
                       if stdout_isatty() {
                           Format::colored()
                       } else {
                           Format::plain()
                       })
}

/// Asynchronos drain to `stderr`
///
/// Automatically using color if output goes to tty
pub fn async_stderr() -> AsyncStreamer<Format> {
    AsyncStreamer::new(io::stderr(),
                       if stderr_isatty() {
                           Format::colored()
                       } else {
                           Format::plain()
                       })
}
<|endoftext|>


<commit_before>extern crate varlink_generator;

use std::fs::File;
use std::io::prelude::*;
// use std::process::Command;

fn compare_files(a: std::path::PathBuf, b: std::path::PathBuf) -> bool {
    let fa = File::open(a).unwrap();
    let fb = File::open(b).unwrap();
    fa.bytes()
        .zip(fb.bytes())
        .all(|(a, b)| a.unwrap() == b.unwrap())
}

#[test]
fn test_generate() {
    std::env::set_var("OUT_DIR", std::env::temp_dir());
    varlink_generator::cargo_build("tests/org.example.complex.varlink");
    let mut path = std::path::PathBuf::from(std::env::temp_dir());
    path.push("org.example.complex.rs");

    // assert!(
    // Command::new("rustfmt")
    // .arg(path.to_str().unwrap())
    // .output()
    // .is_ok()
    // );
    //
    if !compare_files(
        std::path::PathBuf::from("tests/org.example.complex.rs_out"),
        path.clone(),
    ) {
        panic!(
            "Generated files differ: '{}' versus '{}'",
            "tests/org.example.complex.rs_out",
            &path.to_str().unwrap()
        );
    }
    let _ = std::fs::remove_file(path);
}
<commit_msg>varlink_generator/tests/generator.rs: only on unix<commit_after>extern crate varlink_generator;

use std::fs::File;
use std::io::prelude::*;
// use std::process::Command;

fn compare_files(a: std::path::PathBuf, b: std::path::PathBuf) -> bool {
    let fa = File::open(a).unwrap();
    let fb = File::open(b).unwrap();
    fa.bytes()
        .zip(fb.bytes())
        .all(|(a, b)| a.unwrap() == b.unwrap())
}

#[cfg(unix)]
#[test]
fn test_generate() {
    std::env::set_var("OUT_DIR", std::env::temp_dir());
    varlink_generator::cargo_build("tests/org.example.complex.varlink");
    let mut path = std::path::PathBuf::from(std::env::temp_dir());
    path.push("org.example.complex.rs");

    // assert!(
    // Command::new("rustfmt")
    // .arg(path.to_str().unwrap())
    // .output()
    // .is_ok()
    // );
    //
    if !compare_files(
        std::path::PathBuf::from("tests/org.example.complex.rs_out"),
        path.clone(),
    ) {
        panic!(
            "Generated files differ: '{}' versus '{}'",
            "tests/org.example.complex.rs_out",
            &path.to_str().unwrap()
        );
    }
    let _ = std::fs::remove_file(path);
}
<|endoftext|>


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

use crate::common::types::{CliCommand, CliResult, CliTypedResult, TransactionOptions};
use aptos_rest_client::Transaction;
use aptos_transaction_builder::aptos_stdlib;
use aptos_types::account_address::AccountAddress;
use async_trait::async_trait;
use clap::Parser;

/// Tool for manipulating stake
///
#[derive(Parser)]
pub enum StakeTool {
    AddStake(AddStake),
    UnlockStake(UnlockStake),
    WithdrawStake(WithdrawStake),
    IncreaseLockup(IncreaseLockup),
    InitializeStakeOwner(InitializeStakeOwner),
    SetOperator(SetOperator),
    SetDelegatedVoter(SetDelegatedVoter),
}

impl StakeTool {
    pub async fn execute(self) -> CliResult {
        use StakeTool::*;
        match self {
            AddStake(tool) => tool.execute_serialized().await,
            UnlockStake(tool) => tool.execute_serialized().await,
            WithdrawStake(tool) => tool.execute_serialized().await,
            IncreaseLockup(tool) => tool.execute_serialized().await,
            InitializeStakeOwner(tool) => tool.execute_serialized().await,
            SetOperator(tool) => tool.execute_serialized().await,
            SetDelegatedVoter(tool) => tool.execute_serialized().await,
        }
    }
}

/// Stake coins for an account to the stake pool
#[derive(Parser)]
pub struct AddStake {
    /// Amount of coins to add to stake
    #[clap(long)]
    pub amount: u64,

    #[clap(flatten)]
    pub(crate) txn_options: TransactionOptions,
}

#[async_trait]
impl CliCommand<Transaction> for AddStake {
    fn command_name(&self) -> &'static str {
        "AddStake"
    }

    async fn execute(mut self) -> CliTypedResult<Transaction> {
        self.txn_options
            .submit_transaction(aptos_stdlib::stake_add_stake(self.amount))
            .await
    }
}

/// Unlock staked coins
///
/// Coins can only be unlocked if they no longer have an applied lockup period
#[derive(Parser)]
pub struct UnlockStake {
    #[clap(flatten)]
    pub(crate) txn_options: TransactionOptions,
    /// Amount of coins to unlock
    #[clap(long)]
    pub amount: u64,
}

#[async_trait]
impl CliCommand<Transaction> for UnlockStake {
    fn command_name(&self) -> &'static str {
        "UnlockStake"
    }

    async fn execute(mut self) -> CliTypedResult<Transaction> {
        self.txn_options
            .submit_transaction(aptos_stdlib::stake_unlock(self.amount))
            .await
    }
}

/// Withdraw all unlocked staked coins
///
/// Before calling `WithdrawStake`, `UnlockStake` must be called first.
#[derive(Parser)]
pub struct WithdrawStake {
    #[clap(flatten)]
    pub(crate) node_op_options: TransactionOptions,
    /// Amount of coins to withdraw
    #[clap(long)]
    pub amount: u64,
}

#[async_trait]
impl CliCommand<Transaction> for WithdrawStake {
    fn command_name(&self) -> &'static str {
        "WithdrawStake"
    }

    async fn execute(mut self) -> CliTypedResult<Transaction> {
        self.node_op_options
            .submit_transaction(aptos_stdlib::stake_withdraw(self.amount))
            .await
    }
}

/// Increase lockup of all staked coins in an account
#[derive(Parser)]
pub struct IncreaseLockup {
    #[clap(flatten)]
    pub(crate) txn_options: TransactionOptions,
}

#[async_trait]
impl CliCommand<Transaction> for IncreaseLockup {
    fn command_name(&self) -> &'static str {
        "IncreaseLockup"
    }

    async fn execute(mut self) -> CliTypedResult<Transaction> {
        self.txn_options
            .submit_transaction(aptos_stdlib::stake_increase_lockup())
            .await
    }
}

/// Register stake owner, to gain capability to delegate
/// operator or voting capability to a different account.
#[derive(Parser)]
pub struct InitializeStakeOwner {
    #[clap(flatten)]
    pub(crate) txn_options: TransactionOptions,
    #[clap(long)]
    pub initial_stake_amount: u64,
    #[clap(long)]
    pub operator_address: Option<AccountAddress>,
    #[clap(long)]
    pub voter_address: Option<AccountAddress>,
}

#[async_trait]
impl CliCommand<Transaction> for InitializeStakeOwner {
    fn command_name(&self) -> &'static str {
        "InitializeStakeOwner"
    }

    async fn execute(mut self) -> CliTypedResult<Transaction> {
        let owner_address = self.txn_options.sender_address()?;
        self.txn_options
            .submit_transaction(aptos_stdlib::stake_initialize_owner_only(
                self.initial_stake_amount,
                self.operator_address.unwrap_or(owner_address),
                self.voter_address.unwrap_or(owner_address),
            ))
            .await
    }
}

/// Delegate operator (running validator node) capability from the stake owner
/// to the given voter address
#[derive(Parser)]
pub struct SetOperator {
    #[clap(flatten)]
    pub(crate) txn_options: TransactionOptions,
    #[clap(long)]
    pub operator_address: AccountAddress,
}

#[async_trait]
impl CliCommand<Transaction> for SetOperator {
    fn command_name(&self) -> &'static str {
        "SetOperator"
    }

    async fn execute(mut self) -> CliTypedResult<Transaction> {
        self.txn_options
            .submit_transaction(aptos_stdlib::stake_set_operator(self.operator_address))
            .await
    }
}

/// Delegate voting capability from the stake owner to the given voter address
#[derive(Parser)]
pub struct SetDelegatedVoter {
    #[clap(flatten)]
    pub(crate) txn_options: TransactionOptions,
    #[clap(long)]
    pub voter_address: AccountAddress,
}

#[async_trait]
impl CliCommand<Transaction> for SetDelegatedVoter {
    fn command_name(&self) -> &'static str {
        "SetDelegatedVoter"
    }

    async fn execute(mut self) -> CliTypedResult<Transaction> {
        self.txn_options
            .submit_transaction(aptos_stdlib::stake_set_delegated_voter(self.voter_address))
            .await
    }
}
<commit_msg>[aptos-cli] Cleanup argument documentation of staking commands<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::common::types::{CliCommand, CliResult, CliTypedResult, TransactionOptions};
use aptos_rest_client::Transaction;
use aptos_transaction_builder::aptos_stdlib;
use aptos_types::account_address::AccountAddress;
use async_trait::async_trait;
use clap::Parser;

/// Tool for manipulating stake
///
#[derive(Parser)]
pub enum StakeTool {
    AddStake(AddStake),
    UnlockStake(UnlockStake),
    WithdrawStake(WithdrawStake),
    IncreaseLockup(IncreaseLockup),
    InitializeStakeOwner(InitializeStakeOwner),
    SetOperator(SetOperator),
    SetDelegatedVoter(SetDelegatedVoter),
}

impl StakeTool {
    pub async fn execute(self) -> CliResult {
        use StakeTool::*;
        match self {
            AddStake(tool) => tool.execute_serialized().await,
            UnlockStake(tool) => tool.execute_serialized().await,
            WithdrawStake(tool) => tool.execute_serialized().await,
            IncreaseLockup(tool) => tool.execute_serialized().await,
            InitializeStakeOwner(tool) => tool.execute_serialized().await,
            SetOperator(tool) => tool.execute_serialized().await,
            SetDelegatedVoter(tool) => tool.execute_serialized().await,
        }
    }
}

/// Stake coins to the stake pool
///
/// This command allows stake pool owners to add coins to their stake.
#[derive(Parser)]
pub struct AddStake {
    /// Amount of coins to add to stake
    #[clap(long)]
    pub amount: u64,

    #[clap(flatten)]
    pub(crate) txn_options: TransactionOptions,
}

#[async_trait]
impl CliCommand<Transaction> for AddStake {
    fn command_name(&self) -> &'static str {
        "AddStake"
    }

    async fn execute(mut self) -> CliTypedResult<Transaction> {
        self.txn_options
            .submit_transaction(aptos_stdlib::stake_add_stake(self.amount))
            .await
    }
}

/// Unlock staked coins
///
/// Coins can only be unlocked if they no longer have an applied lockup period
#[derive(Parser)]
pub struct UnlockStake {
    /// Amount of coins to unlock
    #[clap(long)]
    pub amount: u64,

    #[clap(flatten)]
    pub(crate) txn_options: TransactionOptions,
}

#[async_trait]
impl CliCommand<Transaction> for UnlockStake {
    fn command_name(&self) -> &'static str {
        "UnlockStake"
    }

    async fn execute(mut self) -> CliTypedResult<Transaction> {
        self.txn_options
            .submit_transaction(aptos_stdlib::stake_unlock(self.amount))
            .await
    }
}

/// Withdraw unlocked staked coins
///
/// This allows users to withdraw stake back into their CoinStore.
/// Before calling `WithdrawStake`, `UnlockStake` must be called first.
#[derive(Parser)]
pub struct WithdrawStake {
    /// Amount of coins to withdraw
    #[clap(long)]
    pub amount: u64,

    #[clap(flatten)]
    pub(crate) node_op_options: TransactionOptions,
}

#[async_trait]
impl CliCommand<Transaction> for WithdrawStake {
    fn command_name(&self) -> &'static str {
        "WithdrawStake"
    }

    async fn execute(mut self) -> CliTypedResult<Transaction> {
        self.node_op_options
            .submit_transaction(aptos_stdlib::stake_withdraw(self.amount))
            .await
    }
}

/// Increase lockup of all staked coins in the stake pool
///
/// Lockup may need to be increased in order to vote on a proposal.
#[derive(Parser)]
pub struct IncreaseLockup {
    #[clap(flatten)]
    pub(crate) txn_options: TransactionOptions,
}

#[async_trait]
impl CliCommand<Transaction> for IncreaseLockup {
    fn command_name(&self) -> &'static str {
        "IncreaseLockup"
    }

    async fn execute(mut self) -> CliTypedResult<Transaction> {
        self.txn_options
            .submit_transaction(aptos_stdlib::stake_increase_lockup())
            .await
    }
}

/// Initialize stake owner
///
/// Initializing stake owner adds the capability to delegate the
/// stake pool to an operator, or delegate voting to a different account.
#[derive(Parser)]
pub struct InitializeStakeOwner {
    /// Initial amount of coins to be staked
    #[clap(long)]
    pub initial_stake_amount: u64,

    /// Account Address of delegated operator
    ///
    /// If not specified, it will be the same as the owner
    #[clap(long)]
    pub operator_address: Option<AccountAddress>,

    /// Account address of delegated voter
    ///
    /// If not specified, it will be the same as the owner
    #[clap(long)]
    pub voter_address: Option<AccountAddress>,

    #[clap(flatten)]
    pub(crate) txn_options: TransactionOptions,
}

#[async_trait]
impl CliCommand<Transaction> for InitializeStakeOwner {
    fn command_name(&self) -> &'static str {
        "InitializeStakeOwner"
    }

    async fn execute(mut self) -> CliTypedResult<Transaction> {
        let owner_address = self.txn_options.sender_address()?;
        self.txn_options
            .submit_transaction(aptos_stdlib::stake_initialize_owner_only(
                self.initial_stake_amount,
                self.operator_address.unwrap_or(owner_address),
                self.voter_address.unwrap_or(owner_address),
            ))
            .await
    }
}

/// Delegate operator capability from the stake owner to another account
#[derive(Parser)]
pub struct SetOperator {
    /// Account Address of delegated operator
    ///
    /// If not specified, it will be the same as the owner
    #[clap(long)]
    pub operator_address: AccountAddress,

    #[clap(flatten)]
    pub(crate) txn_options: TransactionOptions,
}

#[async_trait]
impl CliCommand<Transaction> for SetOperator {
    fn command_name(&self) -> &'static str {
        "SetOperator"
    }

    async fn execute(mut self) -> CliTypedResult<Transaction> {
        self.txn_options
            .submit_transaction(aptos_stdlib::stake_set_operator(self.operator_address))
            .await
    }
}

/// Delegate voting capability from the stake owner to another account
#[derive(Parser)]
pub struct SetDelegatedVoter {
    /// Account Address of delegated voter
    ///
    /// If not specified, it will be the same as the owner
    #[clap(long)]
    pub voter_address: AccountAddress,

    #[clap(flatten)]
    pub(crate) txn_options: TransactionOptions,
}

#[async_trait]
impl CliCommand<Transaction> for SetDelegatedVoter {
    fn command_name(&self) -> &'static str {
        "SetDelegatedVoter"
    }

    async fn execute(mut self) -> CliTypedResult<Transaction> {
        self.txn_options
            .submit_transaction(aptos_stdlib::stake_set_delegated_voter(self.voter_address))
            .await
    }
}
<|endoftext|>