rocket/trace/subscriber/
visit.rs

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
use std::fmt;

use tinyvec::TinyVec;
use tracing::field::{Field, Visit};
use tracing_subscriber::field::RecordFields;

use crate::util::Formatter;

pub trait RecordDisplay: RecordFields {
    fn find_map_display<T, F: Fn(&dyn fmt::Display) -> T>(&self, name: &str, f: F) -> Option<T>;
    fn record_display<F: FnMut(&Field, &dyn fmt::Display)>(&self, f: F);
}

#[derive(Debug)]
pub struct Data {
    // start: Instant,
    map: TinyVec<[(&'static str, String); 3]>,
}

impl Data {
    pub fn new<T: RecordFields>(attrs: T) -> Self {
        let mut data = Data {
            // start: Instant::now(),
            map: TinyVec::new(),
        };

        attrs.record(&mut data);
        data
    }

    pub fn get(&self, key: &str) -> Option<&str> {
        self.map.iter()
            .find(|(k, _)| k == &key)
            .map(|(_, v)| v.as_str())
    }
}

impl std::ops::Index<&str> for Data {
    type Output = str;

    fn index(&self, index: &str) -> &Self::Output {
        self.get(index).unwrap_or("[internal error: missing key]")
    }
}

impl Visit for Data {
    fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {
        self.map.push((field.name(), format!("{:?}", value)));
    }

    fn record_str(&mut self, field: &Field, value: &str) {
        self.map.push((field.name(), value.into()));
    }
}

impl<T: RecordFields> RecordDisplay for T {
    fn find_map_display<V, F: Fn(&dyn fmt::Display) -> V>(&self, name: &str, f: F) -> Option<V> {
        let mut value = None;
        self.record_display(|field, item| if field.name() == name { value = Some(f(item)); });
        value
    }

    fn record_display<F: FnMut(&Field, &dyn fmt::Display)>(&self, f: F) {
        struct DisplayVisit<F>(F);

        impl<F: FnMut(&Field, &dyn fmt::Display)> Visit for DisplayVisit<F> {
            fn record_debug(&mut self, field: &Field, value: &dyn fmt::Debug) {
                (self.0)(field, &Formatter(|f| value.fmt(f)));
            }

            fn record_str(&mut self, field: &Field, value: &str) {
                (self.0)(field, &value)
            }
        }

        self.record(&mut DisplayVisit(f));
    }
}