prefix_trie/rkyv/
serialize.rs

1//! Serialization logic for rkyv
2
3use std::{
4    collections::{HashSet, VecDeque},
5    marker::PhantomData,
6};
7
8use crate::{
9    allocator::Loc,
10    joint::{JointPrefix, JointPrefixMap, JointPrefixSet},
11    rkyv::{
12        map::{MyPhantomData, NodeRepr},
13        ArchivedJointPrefixMap, ArchivedJointPrefixSet, ArchivedPrefixMap, ArchivedPrefixSet,
14        PrefixMapResolver,
15    },
16    table::Table,
17    Prefix, PrefixMap, PrefixSet,
18};
19use rkyv::{
20    munge::munge,
21    rancor::Fallible,
22    ser::{Allocator, Writer},
23    vec::ArchivedVec,
24    Archive, Place, Serialize,
25};
26
27use super::JointPrefixMapResolver;
28
29impl<P: Prefix, T: Archive> Archive for PrefixMap<P, T> {
30    type Archived = ArchivedPrefixMap<P, T>;
31    type Resolver = PrefixMapResolver;
32    fn resolve(&self, resolver: Self::Resolver, out: Place<Self::Archived>) {
33        munge!(let ArchivedPrefixMap { nodes, data, marker } = out);
34        ArchivedVec::resolve_from_len(resolver.nodes_len, resolver.nodes, nodes);
35        ArchivedVec::resolve_from_len(resolver.data_len, resolver.data, data);
36        marker.write(MyPhantomData(PhantomData));
37    }
38}
39
40impl<P: Prefix> Archive for PrefixSet<P> {
41    type Archived = ArchivedPrefixSet<P>;
42    type Resolver = PrefixMapResolver;
43    fn resolve(&self, resolver: Self::Resolver, out: Place<Self::Archived>) {
44        munge!(let ArchivedPrefixSet(out) = out);
45        self.0.resolve(resolver, out)
46    }
47}
48
49impl<P: JointPrefix, T: Archive> Archive for JointPrefixMap<P, T> {
50    type Archived = ArchivedJointPrefixMap<P, T>;
51    type Resolver = JointPrefixMapResolver;
52    fn resolve(&self, resolver: Self::Resolver, out: Place<Self::Archived>) {
53        munge!(let ArchivedJointPrefixMap {t1, t2} = out);
54        self.t1.resolve(resolver.t1, t1);
55        self.t2.resolve(resolver.t2, t2);
56    }
57}
58
59impl<P: JointPrefix> Archive for JointPrefixSet<P> {
60    type Archived = ArchivedJointPrefixSet<P>;
61    type Resolver = JointPrefixMapResolver;
62    fn resolve(&self, resolver: Self::Resolver, out: Place<Self::Archived>) {
63        munge!(let ArchivedJointPrefixSet {t1, t2} = out);
64        self.t1.resolve(resolver.t1, t1);
65        self.t2.resolve(resolver.t2, t2);
66    }
67}
68
69impl<P, T, S> Serialize<S> for PrefixMap<P, T>
70where
71    P: Prefix,
72    T: Serialize<S>,
73    S: Fallible + Writer + Allocator + ?Sized,
74{
75    fn serialize(&self, s: &mut S) -> Result<PrefixMapResolver, S::Error> {
76        let table = self.table();
77
78        // pass 1: Identify empty nodes that do not need to be serialized
79        let mut ignore = HashSet::new();
80        fn find_empty<T>(table: &Table<T>, ignore: &mut HashSet<Loc>, loc: Loc) -> bool {
81            let node = table.node(loc);
82            let mut children_have_value = false;
83            for child in node.child_locs() {
84                children_have_value |= find_empty(table, ignore, child);
85            }
86
87            if loc.is_root() || children_have_value || node.data_bitmap != 0 {
88                true
89            } else {
90                ignore.insert(loc);
91                false
92            }
93        }
94        find_empty(table, &mut ignore, Loc::root());
95
96        // pass 2: the actual serialization
97        let mut nodes = Vec::<NodeRepr>::new();
98        let mut data = Vec::<&T>::with_capacity(self.count);
99        let mut queue = VecDeque::new();
100        nodes.push(NodeRepr::default());
101        queue.push_back(Loc::root());
102
103        let mut cur_node = 0;
104
105        while let Some(loc) = queue.pop_front() {
106            // ignore empty nodes
107            if ignore.contains(&loc) {
108                continue;
109            }
110
111            // update the node info
112            let orig = table.node(loc);
113            nodes[cur_node].data_bitmap = orig.data_bitmap;
114            nodes[cur_node].child_bitmap = orig.child_bitmap;
115            nodes[cur_node].data_idx = data.len() as u32;
116            nodes[cur_node].children_idx = nodes.len() as u32;
117
118            for child_loc in orig.child_locs() {
119                // extend the nodes vector
120                nodes.push(NodeRepr::default());
121                // extend the queue
122                queue.push_back(child_loc);
123            }
124
125            for data_loc in orig.data_locs() {
126                // Safety: we get only live references from data_locs that are not modified since.
127                data.push(unsafe { table.cell(data_loc) })
128            }
129
130            cur_node += 1;
131        }
132
133        debug_assert_eq!(cur_node, nodes.len());
134        debug_assert_eq!(data.len(), self.count);
135
136        let nodes_len = nodes.len();
137        let data_len = data.len();
138        let nodes = ArchivedVec::serialize_from_slice(&nodes, s)?;
139        let data = ArchivedVec::serialize_from_iter::<T, _, _>(data.iter().copied(), s)?;
140
141        Ok(PrefixMapResolver {
142            nodes,
143            nodes_len,
144            data,
145            data_len,
146        })
147    }
148}
149
150impl<P, S> Serialize<S> for PrefixSet<P>
151where
152    P: Prefix,
153    S: Fallible + Writer + Allocator + ?Sized,
154{
155    fn serialize(&self, s: &mut S) -> Result<PrefixMapResolver, S::Error> {
156        self.0.serialize(s)
157    }
158}
159
160impl<P, T, S> Serialize<S> for JointPrefixMap<P, T>
161where
162    P: JointPrefix,
163    T: Archive + Serialize<S>,
164    S: Fallible + Writer + Allocator + ?Sized,
165{
166    fn serialize(&self, s: &mut S) -> Result<JointPrefixMapResolver, S::Error> {
167        let t1 = self.t1.serialize(s)?;
168        let t2 = self.t2.serialize(s)?;
169        Ok(JointPrefixMapResolver { t1, t2 })
170    }
171}
172
173impl<P, S> Serialize<S> for JointPrefixSet<P>
174where
175    P: JointPrefix,
176    S: Fallible + Writer + Allocator + ?Sized,
177{
178    fn serialize(&self, s: &mut S) -> Result<JointPrefixMapResolver, S::Error> {
179        let t1 = self.t1.serialize(s)?;
180        let t2 = self.t2.serialize(s)?;
181        Ok(JointPrefixMapResolver { t1, t2 })
182    }
183}
prefix_trie/rkyv/serialize.rs

prefix_trie/rkyv/
serialize.rs
