osm_io/osm/pbf/
thread_local_accumulator.rs

1use std::cell::RefCell;
2use std::collections::HashMap;
3
4use uuid::Uuid;
5
6use crate::osm::model::element::Element;
7
8thread_local! {
9    static ACCUMULATORS: RefCell<HashMap<String, Vec<Element>>> = RefCell::new(HashMap::new());
10}
11
12/// Accumulate Elements to avoid calling [ParallelWriter] for each element
13pub struct ThreadLocalAccumulator {
14    id: String,
15    capacity: usize,
16}
17
18impl ThreadLocalAccumulator {
19    pub fn new(capacity: usize) -> ThreadLocalAccumulator {
20        let id = Uuid::new_v4().to_string();
21        ThreadLocalAccumulator {
22            id,
23            capacity,
24        }
25    }
26
27    pub fn add(&self, element: Element) {
28        ACCUMULATORS.with(|accumulators| {
29            let mut accumulators = accumulators.borrow_mut();
30            let accumulator = accumulators.get_mut(self.id.as_str());
31            match accumulator {
32                None => {
33                    let mut acc = Vec::with_capacity(self.capacity);
34                    acc.push(element);
35                    accumulators.insert(self.id.clone(), acc);
36                }
37                Some(acc) => {
38                    acc.push(element);
39                }
40            }
41        });
42    }
43
44    pub fn elements(&self) -> Vec<Element> {
45        ACCUMULATORS.with(|accumulators| {
46            if !accumulators.borrow().contains_key(self.id.as_str()) {
47                Vec::new()
48            } else {
49                let mut accumulators = accumulators.borrow_mut();
50                let accumulator = accumulators.get_mut(self.id.as_str()).unwrap();
51                std::mem::replace(accumulator, Vec::with_capacity(self.capacity))
52            }
53        })
54    }
55
56    #[allow(dead_code)]
57    pub(crate) fn len(&self) -> usize {
58        ACCUMULATORS.with(|accumulators| {
59            if !accumulators.borrow().contains_key(self.id.as_str()) {
60                0
61            } else {
62                let accumulators = accumulators.borrow_mut();
63                let accumulator = accumulators.get(self.id.as_str()).unwrap();
64                accumulator.len()
65            }
66        })
67    }
68}
69
osm_io/osm/pbf/thread_local_accumulator.rs

osm_io/osm/pbf/
thread_local_accumulator.rs
