1use crate::prelude::*;
2use std::collections::HashMap;
3use std::default::Default;
4use std::fmt;
5
6#[derive(Default)]
7struct Path {
8 names: String,
9 types: String,
10}
11
12impl Path {
13 fn c(s: &String, x: &impl fmt::Display) -> String {
14 let x = format!("{}", x);
15 if s.is_empty() {
16 x
17 } else {
18 if x.is_empty() {
19 s.clone()
20 } else {
21 format!("{}.{}", s, x)
22 }
23 }
24 }
25
26 #[must_use]
27 pub fn a(&self, name: &impl fmt::Display, type_id: &impl fmt::Display) -> Self {
28 let names = Self::c(&self.names, name);
29 let types = Self::c(&self.types, type_id);
30 Self { names, types }
31 }
32}
33
34struct PathAggregation {
35 types: String,
36 size: usize,
37}
38
39#[derive(Default, Clone)]
40struct TypeAggregation {
41 size: usize,
42 count: usize,
43}
44
45struct SizeBreakdown {
46 by_path: HashMap<String, PathAggregation>,
47 by_type: HashMap<String, TypeAggregation>,
48 total: usize,
49}
50
51impl fmt::Display for SizeBreakdown {
52 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
53 let mut by_path: Vec<_> = self.by_path.iter().collect();
54 let mut by_type: Vec<_> = self.by_type.iter().collect();
55
56 by_path.sort_by_key(|i| usize::MAX - i.1.size);
57 by_type.sort_by_key(|i| usize::MAX - i.1.size);
58
59 writeln!(f, "Largest by path:")?;
60 for (path, agg) in by_path.iter() {
61 writeln!(f, "\t{}\n\t {}\n\t {}", agg.size, path, agg.types)?;
62 }
63
64 writeln!(f)?;
65 writeln!(f, "Largest by type:")?;
66 for (t, agg) in by_type.iter() {
67 writeln!(f, "\t {}x {} @ {}", agg.count, agg.size, t)?;
68 }
69
70 let accounted: usize = by_type.iter().map(|i| (i.1).size).sum();
71
72 writeln!(f)?;
73 writeln!(f, "Other: {}", self.total - accounted)?;
74 writeln!(f, "Total: {}", self.total)?;
75
76 Ok(())
77 }
78}
79
80impl SizeBreakdown {
81 fn add(&mut self, path: &Path, type_id: &'static str, bytes: &Bytes<'_>) {
82 let len = bytes.len();
83 let before = self.by_type.get(type_id).cloned().unwrap_or_default();
84 self.by_type.insert(
85 type_id.to_owned(),
86 TypeAggregation {
87 count: before.count + 1,
88 size: before.size + len,
89 },
90 );
91
92 let types = Path::c(&path.types, &type_id);
93
94 let prev = self.by_path.insert(path.names.clone(), PathAggregation { types, size: len });
95 assert!(prev.is_none());
96 }
97}
98
99fn visit_array(path: Path, branch: &DynArrayBranch, breakdown: &mut SizeBreakdown) {
100 match branch {
101 DynArrayBranch::ArrayFixed { values, len } => visit_array(path.a(&format!("[{}]", len), &"Array Fixed"), values, breakdown),
102 DynArrayBranch::Array { len, values } => {
103 visit_array(path.a(&"len", &"Array"), len, breakdown);
104 visit_array(path.a(&"values", &"Array"), values, breakdown);
105 }
106 DynArrayBranch::Enum { discriminants, variants } => {
107 visit_array(path.a(&"discriminants", &"Enum"), discriminants, breakdown);
108 for variant in variants.iter() {
109 visit_array(path.a(&variant.ident, &"Enum"), &variant.data, breakdown);
110 }
111 }
112 DynArrayBranch::Boolean(enc) => match enc {
113 ArrayBool::Packed(b) => breakdown.add(&path, "Packed Boolean", b),
114 ArrayBool::RLE(_first, runs) => visit_array(path.a(&"runs", &"Bool RLE"), runs, breakdown),
115 },
116 DynArrayBranch::Float(f) => match f {
117 ArrayFloat::DoubleGorilla(b) => breakdown.add(&path, "Gorilla", b),
118 ArrayFloat::F32(b) => breakdown.add(&path, "Fixed F32", b),
119 ArrayFloat::F64(b) => breakdown.add(&path, "Fixed F64", b),
120 ArrayFloat::Zfp32(b) => breakdown.add(&path, "Zfp 64", b),
121 ArrayFloat::Zfp64(b) => breakdown.add(&path, "Zfp 32", b),
122 },
123 DynArrayBranch::Integer(ArrayInteger { bytes, encoding }) => match encoding {
124 ArrayIntegerEncoding::PrefixVarInt => breakdown.add(&path, "Prefix Varint", bytes),
125 ArrayIntegerEncoding::Simple16 => breakdown.add(&path, "Simple16", bytes),
126 ArrayIntegerEncoding::U8 => breakdown.add(&path, "U8 Fixed", bytes),
127 ArrayIntegerEncoding::DeltaZig => breakdown.add(&path, "DeltaZig", bytes),
128 },
129 DynArrayBranch::Map { len, keys, values } => {
130 visit_array(path.a(&"len", &"Map"), len, breakdown);
131 visit_array(path.a(&"keys", &"Map"), keys, breakdown);
132 visit_array(path.a(&"values", &"Map"), values, breakdown);
133 }
134 DynArrayBranch::Object { fields } => {
135 for (name, field) in fields {
136 visit_array(path.a(name, &"Object"), field, breakdown);
137 }
138 }
139 DynArrayBranch::RLE { runs, values } => {
140 visit_array(path.a(&"runs", &"RLE"), runs, breakdown);
141 visit_array(path.a(&"values", &"RLE"), values, breakdown);
142 }
143 DynArrayBranch::Dictionary { indices, values } => {
144 visit_array(path.a(&"indices", &"Dictionary"), indices, breakdown);
145 visit_array(path.a(&"values", &"Dictionary"), values, breakdown);
146 }
147 DynArrayBranch::String(b) => breakdown.add(&path, "UTF-8", b),
148 DynArrayBranch::Tuple { fields } => {
149 for (i, field) in fields.iter().enumerate() {
150 visit_array(path.a(&i, &"Tuple"), field, breakdown);
151 }
152 }
153 DynArrayBranch::Nullable { opt, values } => {
154 visit_array(path.a(&"opt", &"Nullable"), opt, breakdown);
155 visit_array(path.a(&"values", &"Nullable"), values, breakdown);
156 }
157 DynArrayBranch::Void | DynArrayBranch::Map0 | DynArrayBranch::Array0 => {}
158 }
159}
160
161fn visit(path: Path, branch: &DynRootBranch<'_>, breakdown: &mut SizeBreakdown) {
162 match branch {
163 DynRootBranch::Object { fields } => {
164 for (name, value) in fields.iter() {
165 visit(path.a(name, &"Object"), value, breakdown);
166 }
167 }
168 DynRootBranch::Enum { discriminant, value } => visit(path.a(discriminant, &"Enum"), value, breakdown),
169 DynRootBranch::Map { len: _, keys, values } => {
170 visit_array(path.a(&"keys", &"Map"), keys, breakdown);
171 visit_array(path.a(&"values", &"Values"), values, breakdown);
172 }
173 DynRootBranch::Tuple { fields } => {
174 for (i, field) in fields.iter().enumerate() {
175 visit(path.a(&i, &"Tuple"), field, breakdown);
176 }
177 }
178 DynRootBranch::Map1 { key, value } => {
179 visit(path.a(&"key", &"Map1"), key, breakdown);
180 visit(path.a(&"value", &"Map1"), value, breakdown);
181 }
182 DynRootBranch::Array { len, values } => visit_array(path.a(&format!("[{}]", len), &"Array"), values, breakdown),
183 DynRootBranch::Array1(item) => visit(path.a(&"1", &"Array1"), item, breakdown),
184 DynRootBranch::Boolean(_)
185 | DynRootBranch::Array0
186 | DynRootBranch::Map0
187 | DynRootBranch::Void
188 | DynRootBranch::Float(_)
189 | DynRootBranch::Integer(_)
190 | DynRootBranch::String(_) => {}
191 }
192}
193
194pub fn size_breakdown(data: &[u8]) -> DecodeResult<String> {
195 let root = decode_root(data)?;
196
197 let mut breakdown = SizeBreakdown {
198 by_path: HashMap::new(),
199 by_type: HashMap::new(),
200 total: data.len(),
201 };
202 visit(Path::default(), &root, &mut breakdown);
203
204 Ok(format!("{}", breakdown))
205}