oxyroot/rtree/tree/
writer.rs

1use crate::rbytes::{RVersioner, WBuffer};
2use crate::rcont::objarray::WriterObjArray;
3use crate::rdict::StreamerInfo;
4use crate::riofs::file::RootFileStreamerInfoContext;
5use crate::rtree::branch::wbranch::WBranch;
6use crate::rtree::tree::base::Tree;
7use crate::rtree::tree::tio_features::TioFeatures;
8use crate::{rbase, rvers, Marshaler, Named, Object, RootFile};
9use log::trace;
10
11/// Write only Rust equivalent of [`TTree`](https://root.cern/doc/master/classTTree.html)
12///
13/// Mainly used to create [crate::Branch] with name with a provider of data.
14///
15/**
16```
17use oxyroot::RootFile;
18use oxyroot::WriterTree;
19let s = "/tmp/simple.root";
20let mut file = RootFile::create(s).expect("Can not create file");
21let mut tree = WriterTree::new("mytree");
22let it = (0..15);
23tree.new_branch("it", it);
24tree.write(&mut file).expect("Can not write tree");
25file.close().expect("Can not close file");
26```
27 */
28
29pub struct WriterTree {
30    tree: Tree<WBranch<Box<dyn Marshaler>>>,
31    callbacks: Vec<Box<dyn FnMut(StateCallBack)>>,
32}
33
34/// Argument for callbacks called before and during writing branches. Callbacks are mainly used to
35/// monitor the writing process and by the derive [WriteToTree macro](../derive.WriteToTree.html).
36#[derive(Debug)]
37pub enum StateCallBack {
38    Before,
39    /// Before writing branches, called once per entry
40    Branch(String),
41    /// Before writing a branch, called once per entry with name of the branch
42    After,
43}
44
45impl WriterTree {
46    /// Create a new WriterTree with a name. The name is used to identify the TTree in the file.
47    pub fn new<S>(name: S) -> Self
48    where
49        S: AsRef<str>,
50    {
51        Self {
52            tree: Tree {
53                named: rbase::Named::default()
54                    .with_name(name.as_ref().to_string())
55                    .with_title(String::new()),
56                weight: 1.0,
57                scan_field: 25,
58                default_entry_offset_len: 1000,
59                max_entries: 1000000000000,
60                max_entry_loop: 1000000000000,
61                auto_save: -300000000,
62                auto_flush: -30000000,
63                estimate: 1000000,
64                branches: Vec::new(),
65                sinfos: Some(RootFileStreamerInfoContext::new()),
66                ..Default::default()
67            },
68            callbacks: Vec::new(),
69        }
70    }
71
72    pub(crate) fn iobits(&self) -> TioFeatures {
73        self.tree.iobits
74    }
75
76    /// Register callback to be called before and after writing branches. See [StateCallBack](enum.StateCallBack.html) and
77    /// [write](#method.write) method for more information.
78    pub fn add_callback<F>(&mut self, f: Box<F>)
79    where
80        F: FnMut(StateCallBack) + 'static,
81    {
82        self.callbacks.push(f);
83    }
84
85    pub(crate) fn add_streamer(&mut self, si: StreamerInfo) {
86        let sis = self.tree.sinfos.as_mut().unwrap();
87        sis.push(si);
88    }
89
90    /// Add a new branch to the tree.
91    ///
92    /// At this point, the iterator is not consumed, it will be by calling the write method.
93    /// The `T` type has to implement [Marshaler](crate::Marshaler) to be able to write to the file.
94    /// Implementation for basic types are provided by oxyroot.
95    ///
96    /// In order to write a custom type, you have to implement the Marshaler trait or use the derive
97    /// [WriteToTree macro](derive.WriteToTree.html).
98    pub fn new_branch<T, S>(&mut self, name: S, provider: impl Iterator<Item = T> + 'static)
99    where
100        T: Marshaler + 'static,
101        S: AsRef<str>,
102    {
103        let it = provider.map(|x| Box::new(x) as Box<dyn Marshaler>);
104        let wbranchwb = WBranch::new::<T>(name.as_ref(), it, self);
105        self.tree.branches.push(wbranchwb);
106    }
107
108    /// Effectively write the tree to the file.
109    ///
110    /// The branches are written until all provided iterator are exhausted. The branches are written all
111    /// at the same time, the method jump from one iterator to another. During the writing, the callback
112    /// registered with [add_callback](#method.add_callback) are called for each new entry:
113    /// - Before writing branches with the argument [StateCallBack::Before](enum.StateCallBack.html)
114    /// - Before writing each branch with the argument [StateCallBack::Branch](enum.StateCallBack.html)
115    pub fn write(&mut self, file: &mut RootFile) -> crate::riofs::Result<()> {
116        let mut branchs_done = self
117            .tree
118            .branches
119            .iter()
120            .map(|_b| false)
121            .collect::<Vec<_>>();
122        let mut branches = std::mem::take(&mut self.tree.branches);
123        loop {
124            let mut tot = 0;
125            let zip = 0;
126            for f in self.callbacks.iter_mut() {
127                f(StateCallBack::Before);
128            }
129            for (b, d) in branches.iter_mut().zip(branchs_done.iter_mut()) {
130                for f in self.callbacks.iter_mut() {
131                    f(StateCallBack::Branch(b.name().to_string()));
132                }
133                match b.write(self, file)? {
134                    None => *d = true,
135                    Some(nbytes) => {
136                        tot += nbytes;
137                    }
138                }
139            }
140
141            self.tree.tot_bytes += tot as i64;
142            self.tree.zip_bytes += zip as i64;
143            if branchs_done.iter().all(|d| *d) {
144                break;
145            }
146            self.tree.entries += 1;
147        }
148        self.tree.branches = branches;
149
150        trace!(";WriterTree.write_all.entries:{:?}", self.tree.entries);
151
152        self.close(file)
153    }
154
155    fn flush(&mut self, file: &mut RootFile) -> crate::riofs::Result<()> {
156        trace!(";WriterTree.flush:{:?}", true);
157        for b in self.tree.branches.iter_mut() {
158            b.flush(file)?;
159        }
160        Ok(())
161    }
162
163    /// Called by file close method.
164    fn close(&mut self, file: &mut RootFile) -> crate::riofs::Result<()> {
165        trace!(";WriterTree.close:{:?}", true);
166        self.flush(file)?;
167
168        // let t: ReaderTree = self.into();
169
170        file.put(self.tree.named.name(), self)?;
171
172        let sis = self.tree.sinfos.take().unwrap();
173
174        for si in sis.list().iter() {
175            file.add_streamer_info(si.clone());
176        }
177
178        // for sis in self.sinfos.take() {}
179
180        Ok(())
181    }
182}
183
184impl Marshaler for WriterTree {
185    fn marshal(&self, w: &mut WBuffer) -> crate::rbytes::Result<i64> {
186        let len = w.len();
187        let beg = w.pos();
188        trace!(
189            ";WriterTree.marshal.a{beg}.auto_flush:{:?}",
190            self.tree.auto_flush
191        );
192        let hdr = w.write_header(self.class(), Self::rversion(self))?;
193
194        trace!(";WriterTree.marshal.a{beg}.pos.before.named:{:?}", w.pos());
195        w.write_object(&self.tree.named)?;
196        trace!(
197            ";WriterTree.marshal.a{beg}.pos.before.attline:{:?}",
198            w.pos()
199        );
200
201        w.write_object(&self.tree.attline)?;
202        trace!(
203            ";WriterTree.marshal.a{beg}.pos.before.attfill:{:?}",
204            w.pos()
205        );
206
207        // trace!(";WriterTree.marshal.buf.value:{:?}", w.p());
208        w.write_object(&self.tree.attfill)?;
209        trace!(
210            ";WriterTree.marshal.a{beg}.pos.before.attmarker:{:?}",
211            w.pos()
212        );
213        // trace!(";WriterTree.marshal.buf.value:{:?}", w.p());
214        w.write_object(&self.tree.attmarker)?;
215
216        w.write_i64(self.tree.entries)?;
217        w.write_i64(self.tree.tot_bytes)?;
218        w.write_i64(self.tree.zip_bytes)?;
219        w.write_i64(self.tree.saved_bytes)?;
220        w.write_i64(self.tree.flushed_bytes)?;
221        w.write_f64(self.tree.weight)?;
222        w.write_i32(self.tree.timer_interval)?;
223        w.write_i32(self.tree.scan_field)?;
224        w.write_i32(self.tree.update)?;
225        w.write_i32(self.tree.default_entry_offset_len)?;
226
227        w.write_i32(self.tree.clusters.ranges.len().try_into()?)?;
228
229        w.write_i64(self.tree.max_entries)?;
230        w.write_i64(self.tree.max_entry_loop)?;
231        w.write_i64(self.tree.max_virtual_size)?;
232        // trace!(";WriterTree.marshal.buf.value:{:?}", w.p());
233        w.write_i64(self.tree.auto_save)?;
234        // trace!(";WriterTree.marshal.buf.value:{:?}", w.p());
235        trace!(";WriterTree.marshal.auto_flush:{:?}", self.tree.auto_flush);
236        trace!(";WriterTree.marshal.auto_save:{:?}", self.tree.auto_save);
237        trace!(";WriterTree.marshal.estimate:{:?}", self.tree.estimate);
238        w.write_i64(self.tree.auto_flush)?;
239        w.write_i64(self.tree.estimate)?;
240
241        w.write_i8(0)?;
242        w.write_array_i64(&self.tree.clusters.ranges)?;
243        w.write_i8(0)?;
244        w.write_array_i64(&self.tree.clusters.sizes)?;
245        w.write_object(&self.tree.iobits)?;
246        trace!(";WriterTree.marshal.buf.value:{:?}", w.p());
247
248        {
249            let mut branches = WriterObjArray::new();
250            // let tbranches = std::mem::take()
251            for b in self.tree.branches.iter() {
252                branches.push(b, std::ptr::addr_of!(*b) as usize);
253            }
254            w.write_object(&branches)?;
255        }
256
257        trace!(";WriterTree.marshal.buf.value:{:?}", &w.p()[len..]);
258        {
259            let mut leaves = WriterObjArray::new();
260            for b in self.tree.branches.iter() {
261                for leaf in b.branch().tbranch().leaves.iter() {
262                    leaves.push(leaf, std::ptr::addr_of!(*leaf) as usize);
263                }
264            }
265
266            w.write_object(&leaves)?;
267        }
268        {
269            w.write_object_nil()?;
270            w.write_object_nil()?;
271            w.write_object_nil()?;
272            w.write_object_nil()?;
273            w.write_object_nil()?;
274            w.write_object_nil()?;
275            w.write_object_nil()?;
276        }
277
278        let ret = w.set_header(hdr)?;
279        trace!(";WriterTree.marshal.buf.value:{:?}", &w.p()[len..]);
280        Ok(ret)
281
282        // trace!(";WriterTree.marshal.buf.value:{:?}", w.p());
283    }
284}
285
286impl Object for WriterTree {
287    fn class(&self) -> &'_ str {
288        "TTree"
289    }
290}
291
292impl Named for WriterTree {
293    fn name(&self) -> &'_ str {
294        self.tree.named.name()
295    }
296
297    fn title(&self) -> &'_ str {
298        self.tree.named.title()
299    }
300}
301
302impl RVersioner for WriterTree {
303    fn rversion(&self) -> i16 {
304        rvers::TREE
305    }
306}
oxyroot/rtree/tree/writer.rs

oxyroot/rtree/tree/
writer.rs
