...

Text file src/github.com/google/flatbuffers/tests/rust_usage_test/bin/flexbuffers_alloc_check.rs

Documentation: github.com/google/flatbuffers/tests/rust_usage_test/bin

     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
    15extern crate flexbuffers;
    16
    17use flexbuffers::*;
    18use std::alloc::{GlobalAlloc, Layout, System};
    19
    20/// We take over the Rust allocator to count allocations. This is super not thread safe.
    21static mut NUM_ALLOCS: usize = 0;
    22fn current_allocs() -> usize {
    23    unsafe { NUM_ALLOCS }
    24}
    25struct TrackingAllocator;
    26unsafe impl GlobalAlloc for TrackingAllocator {
    27    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
    28        NUM_ALLOCS += 1;
    29        System.alloc(layout)
    30    }
    31    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
    32        System.dealloc(ptr, layout)
    33    }
    34}
    35#[global_allocator]
    36static T: TrackingAllocator = TrackingAllocator;
    37
    38/// Make some example data
    39fn 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.
    70fn 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.
   111fn 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.
   142fn no_extra_allocations() {
   143    main()
   144}

View as plain text