Skip to main content

earley_omnitool/
memory_use.rs

1use std::mem;
2
3use bit_matrix::BitMatrix;
4use bit_vec::BitVec;
5
6use forest::node_handle::NodeHandle;
7use forest::{Bocage, CompactBocage, Forest, NullForest};
8use grammar::InternalGrammar;
9use item::{CompletedItem, Item};
10use recognizer::Recognizer;
11
12const ITEMS_PER_SET: usize = 16;
13
14pub trait MemoryUse {
15    type Arg;
16
17    fn memory_use(&self) -> usize;
18    fn new_with_limit(arg: Self::Arg, memory_limit: usize) -> Self;
19}
20
21impl<'g, F> MemoryUse for Recognizer<'g, F>
22where
23    F: MemoryUse<Arg = &'g InternalGrammar> + Forest,
24{
25    type Arg = &'g InternalGrammar;
26
27    fn memory_use(&self) -> usize {
28        self.forest.memory_use()
29            + self.predicted.memory_use()
30            + self.medial.memory_use()
31            + self.complete.memory_use()
32            + self.indices.memory_use()
33    }
34
35    fn new_with_limit(grammar: &'g InternalGrammar, memory_limit: usize) -> Self {
36        let forest_use_bytes = memory_limit * F::FOREST_BYTES_PER_RECOGNIZER_BYTE
37            / (F::FOREST_BYTES_PER_RECOGNIZER_BYTE + 1);
38        let complete_use = match memory_limit {
39            0..=1000 => 16,
40            1000..=500_000 => 32,
41            500_000..=2_000_000 => 64,
42            _ => 128,
43        };
44        let recognizer_use_bytes = memory_limit
45            - forest_use_bytes
46            - complete_use * mem::size_of::<CompletedItem<F::NodeRef>>();
47        let bytes_per_set = mem::size_of::<usize>()
48            + (grammar.num_syms() + 31) / 32 * 4
49            + ITEMS_PER_SET * mem::size_of::<Item<F::NodeRef>>();
50        let sets_use = recognizer_use_bytes / bytes_per_set;
51        let mut recognizer = Recognizer {
52            forest: F::new_with_limit(grammar, forest_use_bytes),
53            grammar,
54            // The initial location is 0.
55            earleme: 0,
56            // The first Earley set begins at 0 and ends at 0. The second Earley set begins at 0.
57            indices: Vec::with_capacity(sets_use),
58            current_medial_start: 0,
59            // Reserve some capacity for vectors.
60            predicted: BitMatrix::new(sets_use, grammar.num_syms()),
61            medial: Vec::with_capacity(sets_use * ITEMS_PER_SET),
62            complete: Vec::with_capacity(complete_use),
63            lookahead_hint: None,
64        };
65        recognizer.indices.push(0);
66        recognizer.indices.push(0);
67        recognizer.predict(grammar.start_sym());
68        recognizer
69    }
70}
71
72impl<'g, F> Recognizer<'g, F>
73where
74    F: MemoryUse<Arg = &'g InternalGrammar> + Forest,
75{
76    #[inline]
77    pub fn new_with_hint(grammar: &'g InternalGrammar, tokens: usize) -> Self {
78        let forest_use_bytes = tokens * 16;
79        let complete_use = match tokens {
80            0..=200 => 16,
81            200..=10_000 => 32,
82            10_000..=100_000 => 64,
83            _ => 128,
84        };
85        let mut recognizer = Recognizer {
86            forest: F::new_with_limit(grammar, forest_use_bytes),
87            grammar,
88            // The initial location is 0.
89            earleme: 0,
90            // The first Earley set begins at 0 and ends at 0. The second Earley set begins at 0.
91            indices: Vec::with_capacity(tokens + 1),
92            current_medial_start: 0,
93            // Reserve some capacity for vectors.
94            predicted: BitMatrix::new(tokens + 1, grammar.num_syms()),
95            medial: Vec::with_capacity(tokens * ITEMS_PER_SET),
96            complete: Vec::with_capacity(complete_use),
97            lookahead_hint: None,
98        };
99        recognizer.indices.push(0);
100        recognizer.indices.push(0);
101        recognizer.predict(grammar.start_sym());
102        recognizer
103    }
104}
105
106impl<'g> MemoryUse for Recognizer<'g, NullForest> {
107    type Arg = &'g InternalGrammar;
108
109    fn memory_use(&self) -> usize {
110        self.forest.memory_use()
111            + self.predicted.memory_use()
112            + self.medial.memory_use()
113            + self.complete.memory_use()
114            + self.indices.memory_use()
115    }
116
117    fn new_with_limit(grammar: &'g InternalGrammar, memory_limit: usize) -> Self {
118        let complete_use = match memory_limit {
119            0..=1000 => 16,
120            1000..=500_000 => 32,
121            500_000..=2_000_000 => 64,
122            _ => 128,
123        };
124        let recognizer_use_bytes =
125            memory_limit - complete_use * mem::size_of::<CompletedItem<()>>();
126        let bytes_per_set = mem::size_of::<usize>()
127            + (grammar.num_syms() + 31) / 32 * 4
128            + ITEMS_PER_SET * mem::size_of::<Item<()>>();
129        let sets_use = recognizer_use_bytes / bytes_per_set;
130        let mut recognizer = Recognizer {
131            forest: NullForest,
132            grammar,
133            // The initial location is 0.
134            earleme: 0,
135            // The first Earley set begins at 0 and ends at 0. The second Earley set begins at 0.
136            indices: Vec::with_capacity(sets_use),
137            current_medial_start: 0,
138            // Reserve some capacity for vectors.
139            predicted: BitMatrix::new(sets_use, grammar.num_syms()),
140            medial: Vec::with_capacity(sets_use * ITEMS_PER_SET),
141            complete: Vec::with_capacity(complete_use),
142            lookahead_hint: None,
143        };
144        recognizer.indices.push(0);
145        recognizer.indices.push(0);
146        recognizer.predict(grammar.start_sym());
147        recognizer
148    }
149}
150
151impl<T> MemoryUse for Vec<T> {
152    type Arg = ();
153
154    fn memory_use(&self) -> usize {
155        self.capacity() * mem::size_of::<T>()
156    }
157
158    fn new_with_limit(_arg: (), memory_limit: usize) -> Self {
159        let capacity = memory_limit / mem::size_of::<T>();
160        Self::with_capacity(capacity)
161    }
162}
163
164impl MemoryUse for BitMatrix {
165    type Arg = usize;
166
167    fn memory_use(&self) -> usize {
168        let (rows, columns) = self.size();
169        rows * ((columns + 31) / 32 * 4)
170    }
171
172    fn new_with_limit(num_columns: usize, memory_limit: usize) -> Self {
173        let row_size = (num_columns + 31) / 32 * 4;
174        let capacity = memory_limit / row_size;
175        Self::new(capacity, num_columns)
176    }
177}
178
179impl MemoryUse for BitVec {
180    type Arg = ();
181
182    fn memory_use(&self) -> usize {
183        (self.capacity() + 31) / 32 * 4
184    }
185
186    fn new_with_limit(_arg: (), memory_limit: usize) -> Self {
187        let capacity = memory_limit * 8;
188        Self::with_capacity(capacity)
189    }
190}
191
192impl MemoryUse for NullForest {
193    type Arg = ();
194
195    fn memory_use(&self) -> usize {
196        0
197    }
198
199    fn new_with_limit(_arg: (), _memory_limit: usize) -> Self {
200        NullForest
201    }
202}
203
204impl<'g> MemoryUse for Bocage<&'g InternalGrammar> {
205    type Arg = &'g InternalGrammar;
206
207    fn memory_use(&self) -> usize {
208        self.graph.memory_use() + self.gc.liveness.memory_use() + self.gc.dfs.memory_use()
209    }
210
211    fn new_with_limit(grammar: &'g InternalGrammar, memory_limit: usize) -> Self {
212        let dfs_size = match memory_limit {
213            0..=1000 => 8,
214            1000..=100_000 => 32,
215            _ => 64,
216        };
217        let remaining_use = memory_limit - dfs_size * std::mem::size_of::<NodeHandle>();
218        let bytes_per_node = mem::size_of::<u16>() as f32 + 1.0 / 8.0;
219        let graph_size = (remaining_use as f32 / bytes_per_node) as usize;
220        Bocage::with_capacities(grammar, graph_size, dfs_size)
221    }
222}
223
224impl<'g> MemoryUse for CompactBocage<&'g InternalGrammar> {
225    type Arg = &'g InternalGrammar;
226
227    fn memory_use(&self) -> usize {
228        self.graph.vec.memory_use() + self.gc.liveness.memory_use() + self.gc.dfs.memory_use()
229    }
230
231    fn new_with_limit(grammar: &'g InternalGrammar, memory_limit: usize) -> Self {
232        let dfs_size = match memory_limit {
233            0..=1000 => 8,
234            1000..=100_000 => 32,
235            _ => 64,
236        };
237        let remaining_use = memory_limit - dfs_size * std::mem::size_of::<NodeHandle>();
238        let bytes_per_node = mem::size_of::<u16>() as f32 + 1.0 / 8.0;
239        let graph_size = (remaining_use as f32 / bytes_per_node) as usize;
240        CompactBocage::with_capacities(grammar, graph_size, dfs_size)
241    }
242}