1 // Copyright 2019 Google LLC
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 // https://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14
15 extern crate flexbuffers;
16
17 use flexbuffers::*;
18 use std::alloc::{GlobalAlloc, Layout, System};
19
20 /// We take over the Rust allocator to count allocations. This is super not thread safe.
21 static mut NUM_ALLOCS: usize = 0;
current_allocs() -> usize22 fn current_allocs() -> usize {
23 unsafe { NUM_ALLOCS }
24 }
25 struct TrackingAllocator;
26 unsafe impl GlobalAlloc for TrackingAllocator {
alloc(&self, layout: Layout) -> *mut u827 unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
28 NUM_ALLOCS += 1;
29 System.alloc(layout)
30 }
dealloc(&self, ptr: *mut u8, layout: Layout)31 unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
32 System.dealloc(ptr, layout)
33 }
34 }
35 #[global_allocator]
36 static T: TrackingAllocator = TrackingAllocator;
37
38 /// Make some example data
make_monster(mut monster: MapBuilder)39 fn make_monster(mut monster: MapBuilder) {
40 monster.push("type", "great orc");
41 monster.push("age", 100u8);
42 monster.push("name", "Mr. Orc");
43 monster.push("coins", &[1, 25, 50, 100, 250]);
44 monster.push("color", &[255u8, 0, 0, 0]);
45 {
46 let mut weapons = monster.start_vector("weapons");
47 {
48 let mut hammer = weapons.start_map();
49 hammer.push("name", "hammer");
50 hammer.push("damage type", "crush");
51 hammer.push("damage", 20);
52 }
53 {
54 let mut axe = weapons.start_map();
55 axe.push("name", "Great Axe");
56 axe.push("damage type", "slash");
57 axe.push("damage", 30);
58 }
59 }
60 {
61 let mut sounds = monster.start_vector("sounds");
62 sounds.push("grr");
63 sounds.push("rawr");
64 sounds.push("muahaha");
65 }
66 // TODO(cneo): Directly pushing string slices has alloc.
67 }
68
69 // Read back the data from make_monster.
validate_monster(flexbuffer: &[u8])70 fn validate_monster(flexbuffer: &[u8]) {
71 let r = Reader::get_root(flexbuffer).unwrap().as_map();
72
73 assert!(!r.is_empty());
74 assert!(r.index_key("not_a_field").is_none());
75
76 assert_eq!(r.len(), 7);
77 assert_eq!(r.idx("type").as_str(), "great orc");
78 assert_eq!(r.idx("age").as_u8(), 100);
79 assert_eq!(r.idx("name").as_str(), "Mr. Orc");
80
81 let coins = r.idx("coins").as_vector();
82 for (i, &c) in [1, 25, 50, 100, 250].iter().enumerate() {
83 assert_eq!(coins.idx(i).as_u16(), c);
84 }
85 let color = r.idx("color").as_vector();
86 for (i, &c) in [255, 0, 0, 0].iter().enumerate() {
87 assert_eq!(color.idx(i).as_i32(), c);
88 }
89 let weapons = r.idx("weapons").as_vector();
90 assert_eq!(weapons.len(), 2);
91
92 let hammer = weapons.idx(0).as_map();
93 assert_eq!(hammer.idx("name").as_str(), "hammer");
94 assert_eq!(hammer.idx("damage type").as_str(), "crush");
95 assert_eq!(hammer.idx("damage").as_u64(), 20);
96
97 let axe = weapons.idx(1).as_map();
98 assert_eq!(axe.idx("name").as_str(), "Great Axe");
99 assert_eq!(axe.idx("damage type").as_str(), "slash");
100 assert_eq!(axe.idx("damage").as_u64(), 30);
101
102 let sounds = r.idx("sounds").as_vector();
103 for (i, &s) in ["grr", "rawr", "muahaha"].iter().enumerate() {
104 assert_eq!(sounds.idx(i).as_str(), s);
105 }
106 }
107
108 // This is in a separate binary than tests because taking over the global allocator is not
109 // hermetic and not thread safe.
110 #[cfg(not(miri))] // slow.
main()111 fn main() {
112 let start_up = current_allocs();
113
114 // Let's build a flexbuffer from a new (cold) flexbuffer builder.
115 let mut builder = Builder::default();
116 make_monster(builder.start_map());
117 let after_warmup = current_allocs();
118
119 // The builder makes some allocations while warming up.
120 assert!(after_warmup > start_up);
121 assert!(after_warmup < start_up + 20);
122
123 // A warm builder should make no allocations.
124 make_monster(builder.start_map());
125 assert_eq!(after_warmup, current_allocs());
126
127 // Nor should a reader.
128 validate_monster(builder.view());
129 assert_eq!(after_warmup, current_allocs());
130
131 // Do it again just for kicks.
132 make_monster(builder.start_map());
133 validate_monster(builder.view());
134 assert_eq!(after_warmup, current_allocs());
135
136 let final_allocs = current_allocs(); // dbg! does allocate.
137 dbg!(start_up, after_warmup, final_allocs);
138 }
139
140 #[test]
141 #[cfg(not(miri))] // slow.
no_extra_allocations()142 fn no_extra_allocations() {
143 main()
144 }
145