oxyroot/rtree/branch/
mod.rs

1mod tbranch;
2mod tbranch_element;
3mod tbranch_props;
4
5pub(crate) mod wbranch;
6
7pub(crate) use crate::rtree::branch::tbranch::TBranch;
8pub(crate) use crate::rtree::branch::tbranch_element::TBranchElement;
9use std::fmt::Debug;
10
11use crate::rbytes::rbuffer::RBuffer;
12use crate::rbytes::UnmarshalerInto;
13use crate::riofs::file::{RootFileReader, RootFileStreamerInfoContext};
14use crate::root::traits::{Named, Object};
15
16use crate::rtree::streamer_type::type_name_cpp_to_rust;
17use crate::rtypes::FactoryItemRead;
18use log::trace;
19use std::marker::PhantomData;
20
21pub(crate) enum BranchChunks {
22    RegularSized((i32, i32, Vec<u8>)),
23    IrregularSized((i32, Vec<Vec<u8>>, i32)), // _,_, header_bytes
24}
25
26/// Rust equivalent of [`TBranch`](https://root.cern/doc/master/classTBranch.html)
27/// or [`TBranchElement`](https://root.cern/doc/master/classTBranchElement.html) (ie column) of a TTree
28///
29/// Choice between `TBranch` or `TBranchElement` is done when Root file is read.
30/// [Branch] should not be constructed by user but accessed via [crate::Tree::branch]
31#[derive(Debug)]
32pub enum Branch {
33    Base(TBranch),
34    Element(TBranchElement),
35}
36
37impl From<Box<dyn FactoryItemRead>> for Branch {
38    fn from(obj: Box<dyn FactoryItemRead>) -> Self {
39        match obj.class() {
40            "TBranch" => Branch::Base(*obj.downcast::<TBranch>().unwrap()),
41            "TBranchElement" => Branch::Element(*obj.downcast::<TBranchElement>().unwrap()),
42            &_ => todo!(),
43        }
44    }
45}
46
47// impl Debug for Branch {
48//     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
49//         match self {
50//             Branch::Base(_) => f.write_str("TBranch{")?,
51//             Branch::Element(_) => f.write_str("TBranchElement{")?,
52//         };
53//         f.write_str(format!("name: {}", self.name()).as_str())?;
54//         f.write_str(format!("item_t: {}", self.item_type_name()).as_str())?;
55//         f.write_str("}")
56//     }
57// }
58
59impl Branch {
60    fn tbranch_mut(&mut self) -> &mut TBranch {
61        match self {
62            Branch::Base(ref mut bb) => bb,
63            Branch::Element(ref mut be) => &mut be.branch,
64        }
65    }
66
67    pub(crate) fn tbranch(&self) -> &TBranch {
68        match self {
69            Branch::Base(bb) => bb,
70            Branch::Element(be) => &be.branch,
71        }
72    }
73
74    pub fn name(&self) -> &str {
75        let b: &TBranch = self.into();
76        b.name()
77    }
78
79    /// Return the class name of the branch, ie TBranch or TBranchElement.
80    pub fn class(&self) -> &str {
81        match &self {
82            Branch::Base(b) => b.class(),
83            Branch::Element(e) => e.class(),
84        }
85    }
86
87    /// C++ type contained in this branch.
88    pub fn item_type_name(&self) -> String {
89        match self {
90            Branch::Base(bb) => bb.item_type_name(),
91            Branch::Element(be) => be.item_type_name(),
92        }
93    }
94
95    /// Rust equivalent of C++ type returned by [`item_type_name`](crate::Branch::item_type_name)
96    pub fn interpretation(&self) -> String {
97        type_name_cpp_to_rust(self.item_type_name().as_str())
98    }
99
100    /// Number of entries
101    pub fn entries(&self) -> i64 {
102        let b: &TBranch = self.into();
103        b.entries()
104    }
105
106    /// Get iterator over top-level branches
107    pub fn branches(&self) -> impl Iterator<Item = &Branch> {
108        match self {
109            Branch::Base(bb) => bb.branches().iter(),
110            Branch::Element(be) => be.branch.branches().iter(),
111        }
112    }
113
114    /// Get all (recursively) branches in this Branch
115    pub fn branches_r(&self) -> Vec<&Branch> {
116        let mut v = Vec::new();
117
118        for b in self.branches() {
119            v.push(b);
120            for bb in b.branches_r() {
121                v.push(bb);
122            }
123        }
124
125        v
126    }
127
128    /// search in children branches
129    pub fn branch(&self, name: &str) -> Option<&Branch> {
130        match self {
131            Branch::Base(bb) => bb.branch(name),
132            Branch::Element(be) => be.branch.branch(name),
133        }
134    }
135
136    pub(crate) fn set_top_level(&mut self, v: Option<bool>) {
137        match self {
138            Branch::Base(bb) => bb.props.is_top_level = v,
139            Branch::Element(be) => be.set_is_top_level(v),
140        }
141    }
142
143    pub(crate) fn set_reader(&mut self, reader: Option<RootFileReader>) {
144        match self {
145            Branch::Base(bb) => bb.set_reader(Some(reader.unwrap())),
146            Branch::Element(be) => be.branch.set_reader(Some(reader.unwrap())),
147        }
148    }
149
150    pub(crate) fn set_streamer_info(&mut self, sinfos: RootFileStreamerInfoContext) {
151        match self {
152            Branch::Base(bb) => bb.set_streamer_info(sinfos),
153            Branch::Element(be) => be.branch.set_streamer_info(sinfos),
154        }
155    }
156
157    fn get_baskets_buffer(&self) -> Box<dyn Iterator<Item = BranchChunks> + '_> {
158        match self {
159            Branch::Base(bb) => bb.get_baskets_buffer(),
160            Branch::Element(be) => be.get_baskets_buffer(),
161        }
162    }
163
164    pub fn get_basket<'a, F, T>(&'a self, mut func: F) -> impl Iterator<Item = T> + 'a
165    where
166        T: 'a,
167        F: FnMut(&mut RBuffer) -> T + 'a,
168    {
169        trace!("get_basket in BRANCH = {}", self.name());
170
171        let tbranch = match self {
172            Branch::Base(bb) => bb,
173            Branch::Element(be) => &be.branch,
174        };
175
176        trace!(
177            "get_basket in TBranch = {}, nb_branches = {} nb entries = {}",
178            tbranch.name(),
179            tbranch.branches().len(),
180            tbranch.entries()
181        );
182        assert!(tbranch.reader().is_some());
183
184        let it = if !tbranch.branches().is_empty() {
185            let b: Box<dyn Iterator<Item = T>> = Box::new(
186                ZiperBranches::<usize>::new(tbranch.branches(), tbranch.entries() as u32).map(
187                    move |(_n, _chunk_size, buf)| {
188                        let mut r = RBuffer::new(&buf, 0);
189                        func(&mut r)
190                    },
191                ),
192            );
193            b
194        } else {
195            let b: Box<dyn Iterator<Item = T>> =
196                Box::new(self.get_baskets_buffer().flat_map(move |chunk| {
197                    match chunk {
198                        BranchChunks::RegularSized((n, _chunk_size, buf)) => {
199                            let mut r = RBuffer::new(&buf, 0);
200                            let mut v = Vec::with_capacity(n as usize);
201
202                            for _i in 0..n {
203                                v.push(func(&mut r));
204                            }
205                            v
206                        }
207                        BranchChunks::IrregularSized((_n, data_chuncked, header_bytes)) => {
208                            trace!(";Branch.get_baskets.unzip.IrregularSized.call:{:?}", true);
209                            trace!(
210                                ";Branch.get_baskets.unzip.IrregularSized.start.header_bytes:{:?}",
211                                header_bytes
212                            );
213
214                            data_chuncked
215                                .iter()
216                                .map(|buf| {
217                                    trace!("buf = {:?}", buf);
218                                    // if buf.is_empty() {
219                                    //     return T::default();
220                                    // }
221
222                                    trace!(
223                                        ";Branch.get_baskets.unzip.IrregularSized.map.buf:{:?}",
224                                        buf
225                                    );
226
227                                    let mut r = RBuffer::new(buf, 0);
228                                    r.set_skip_header(Some(header_bytes));
229
230                                    func(&mut r)
231                                })
232                                .collect::<Vec<_>>()
233                        }
234                    }
235                }));
236            b
237        };
238
239        it
240    }
241
242    /// Create an iterator over the data of a column (`TBranch`)
243    pub fn as_iter<'a, T>(&'a self) -> crate::Result<impl Iterator<Item = T> + 'a>
244    where
245        T: UnmarshalerInto<Item = T> + 'a,
246    {
247        // println!("typename of branch: {}", self.item_type_name());
248        // println!("typename of type: {:?}", T::classe_name());
249        // println!("typename of type: {:?}", type_name::<T>());
250
251        let ok_typename = match T::classe_name() {
252            None => true,
253            Some(tys) => tys.contains(&self.item_type_name()),
254        };
255
256        if !ok_typename {
257            Err(crate::error::Error::TypeMismatch {
258                given: format!("one of {:?}", T::classe_name().unwrap()),
259                expected: self.item_type_name(),
260            })
261        } else {
262            Ok(self.get_basket(|r| r.read_object_into::<T>().unwrap()))
263        }
264    }
265
266    /// Create an iterator over the data of a column (`TBranch`) but will not check if the provided `T` is
267    /// compatible with actual C++ type of the column.
268    pub fn as_iter_unchecked<'a, T>(&'a self) -> impl Iterator<Item = T> + 'a
269    where
270        T: UnmarshalerInto<Item = T> + 'a,
271    {
272        self.get_basket(|r| r.read_object_into::<T>().unwrap())
273    }
274
275    pub(crate) fn _streamer_type(&self) -> Option<i32> {
276        match self {
277            Branch::Base(_bb) => None,
278            Branch::Element(be) => Some(be.streamer_type()),
279        }
280    }
281}
282
283pub struct ZiperBranchInnerO<'a, T> {
284    pub num_entries: u32,
285    pub chunk_size: i32,
286    pub i: Vec<u8>,
287    // pub o: &'a [u8],
288    phantom: PhantomData<&'a T>,
289}
290
291impl<'a, T> Iterator for ZiperBranchInnerO<'a, T> {
292    type Item = (u32, i32, Vec<u8>);
293
294    fn next(&mut self) -> Option<Self::Item> {
295        // let output_size =
296
297        let o = self.i[0..self.chunk_size as usize].to_vec();
298
299        trace!("o = {:?}", o);
300
301        Some((self.num_entries, self.chunk_size, o))
302    }
303}
304
305pub struct ZiperBranches<'a, T> {
306    _branches: &'a Vec<Branch>,
307    phantom: PhantomData<T>,
308    iterators: Vec<Box<dyn Iterator<Item = BranchChunks> + 'a>>,
309    // output_buffers: Option<ZiperBranchInnerO<'a, T>>,
310    output_buffers: Vec<Option<BranchChunks>>,
311    current_size: Vec<usize>,
312    nb_entries: Vec<i32>,
313}
314
315impl<'a, T> ZiperBranches<'a, T> {
316    pub fn new(branches: &'a Vec<Branch>, _nb_entries: u32) -> Self {
317        let mut v = Vec::new();
318        // let mut v: Vec<dyn Iterator<Item = (u32, i32, Vec<u8>)>> = Vec::new();
319        for branch in branches {
320            let tbranch: &TBranch = branch.into();
321            let data = tbranch.get_baskets_buffer();
322
323            // let b = data.collect::<Vec<_>>();
324            //
325            // trace!("b = {:?}", b);
326
327            // todo!();
328
329            v.push(data);
330
331            // let d = data as &dyn Iterator<Item = (u32, i32, Vec<u8>)>;
332            // v.push(Box::new(data));
333        }
334
335        // todo!();
336
337        ZiperBranches {
338            _branches: branches,
339            phantom: Default::default(),
340            iterators: v,
341            output_buffers: Vec::new(),
342            current_size: Default::default(),
343            nb_entries: Default::default(),
344        }
345    }
346}
347
348impl<'a, T> ZiperBranches<'a, T> {
349    fn fill_output(&mut self) {
350        for it_branch in &mut self.iterators {
351            if let Some(chunk) = it_branch.next() {
352                // let (n, chunk_size, buf) = data;
353                // trace!("n = {}", n);
354
355                let n = match chunk {
356                    BranchChunks::RegularSized((n, _, _)) => n,
357                    BranchChunks::IrregularSized((n, _, _)) => n,
358                };
359
360                self.output_buffers.push(Some(chunk));
361                self.nb_entries.push(n);
362                self.current_size.push(0);
363            }
364        }
365    }
366
367    fn fill_output_one_branch(&mut self, num_branch: usize) {
368        trace!("self.fill_output_one_branch, num_branch = {}", num_branch);
369
370        let it_branch = &mut self.iterators[num_branch];
371        if let Some(chunk) = it_branch.next() {
372            let n = match chunk {
373                BranchChunks::RegularSized((n, _, _)) => n,
374                BranchChunks::IrregularSized((n, _, _)) => n,
375            };
376
377            trace!("n = {}", n);
378
379            self.output_buffers[num_branch] = Some(chunk);
380            self.nb_entries[num_branch] = n;
381            self.current_size[num_branch] = 0;
382        }
383    }
384}
385
386impl<'a, T> Iterator for ZiperBranches<'a, T> {
387    type Item = (u32, i32, Vec<u8>);
388
389    fn next(&mut self) -> Option<Self::Item> {
390        // let mut ret = Vec::new();
391
392        if self.output_buffers.is_empty() {
393            trace!("self.output_buffers.is_empty()");
394            self.fill_output();
395        }
396
397        if self.output_buffers.is_empty() {
398            return None;
399        }
400
401        for b in &self.output_buffers {
402            if b.is_none() {
403                return None;
404            }
405        }
406
407        let size = self.output_buffers.iter().fold(0_usize, |acc, par| {
408            let s = match par.as_ref().unwrap() {
409                BranchChunks::RegularSized((_, s, _)) => s,
410                BranchChunks::IrregularSized(_) => {
411                    todo!()
412                }
413            };
414            acc + *s as usize
415        });
416        // let mut outbuf = vec![0 as u8; size];
417
418        let mut outbuf: Vec<u8> = Vec::with_capacity(size);
419
420        // for (ib, bbuffer) in enumerate(&self.output_buffers) {
421        for ib in 0..self.current_size.len() {
422            match &self.output_buffers[ib] {
423                None => {
424                    panic!("faut remplit");
425                }
426                Some(chunk) => {
427                    let (chunk_size, buf) = match chunk {
428                        BranchChunks::RegularSized((_, c, b)) => (c, b),
429                        BranchChunks::IrregularSized(_) => {
430                            todo!()
431                        }
432                    };
433
434                    let csize = *chunk_size as usize;
435                    let begin = self.current_size[ib] * csize;
436                    let end = (self.current_size[ib] + 1) * csize;
437
438                    let mut ibuffer = buf[begin..end].to_vec();
439                    outbuf.append(&mut ibuffer);
440
441                    self.current_size[ib] += 1;
442
443                    // trace!(
444                    //     "self.current_size = {:?}, self.nb_entries = {:?}",
445                    //     self.current_size,
446                    //     self.nb_entries
447                    // );
448
449                    if self.current_size[ib] == self.nb_entries[ib] as usize {
450                        self.output_buffers[ib] = None;
451
452                        self.fill_output_one_branch(ib);
453                    }
454                }
455            }
456        }
457
458        Some((0, size as i32, outbuf))
459    }
460}
oxyroot/rtree/branch/mod.rs

oxyroot/rtree/branch/
mod.rs
