1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98

use rustc_data_structures::sync::Lrc;
use syntax::ast::*;
use syntax::ptr::P;
use syntax::symbol::Symbol;

use crate::ast_manip::MutVisitNodes;
use crate::command::{CommandState, Registry};
use crate::transform::Transform;
use crate::RefactorCtxt;


/// # `bytestr_to_str` Command
/// 
/// Usage: `bytestr_to_str`
/// 
/// Marks: `target`
/// 
/// Convert bytestring literal expressions marked `target` to string literal
/// expressions.
/// 
/// Note the mark must be placed on the expression, as it is currently difficult to
/// mark a literal node.
pub struct ByteStrToStr;

impl Transform for ByteStrToStr {
    fn transform(&self, krate: &mut Crate, st: &CommandState, _cx: &RefactorCtxt) {
        MutVisitNodes::visit(krate, |e: &mut P<Expr>| {
            if !st.marked(e.id, "target") {
                return;
            }

            match &mut e.node {
                ExprKind::Lit(l) => {
                    match l.node {
                        LitKind::ByteStr(ref bs) => {
                            let s = String::from_utf8((**bs).clone()).unwrap();
                            l.node = LitKind::Str(Symbol::intern(&s), StrStyle::Cooked)
                        }
                        _ => {}
                    }
                }
                _ => {}
            }
        })
    }
}


/// # `remove_null_terminator` Command
/// 
/// Usage: `remove_null_terminator`
/// 
/// Marks: `target`
/// 
/// Remove a trailing `\0` character from marked string and bytestring literal
/// expressions.
/// 
/// Note the mark must be placed on the expression, as it is currently difficult to
/// mark a literal node.
pub struct RemoveNullTerminator;

impl Transform for RemoveNullTerminator {
    fn transform(&self, krate: &mut Crate, st: &CommandState, _cx: &RefactorCtxt) {
        MutVisitNodes::visit(krate, |e: &mut P<Expr>| {
            if !st.marked(e.id, "target") {
                return;
            }

            match &mut e.node {
                ExprKind::Lit(l) => {
                    match &mut l.node {
                        LitKind::ByteStr(bs) => {
                            if bs.last() == Some(&0) {
                                Lrc::get_mut(bs).unwrap().pop();
                            }
                        }
                        LitKind::Str(ref mut s, _style) => {
                            if s.as_str().ends_with("\0") {
                                let end = s.as_str().len() - 1;
                                *s = Symbol::intern(&s.as_str()[..end]);
                            }
                        }
                        _ => {}
                    }
                }
                _ => {}
            }
        });
    }
}


pub fn register_commands(reg: &mut Registry) {
    use super::mk;
    reg.register("bytestr_to_str", |_args| mk(ByteStrToStr));
    reg.register("remove_null_terminator", |_args| mk(RemoveNullTerminator));
}