Skip to main content

fcb_core/writer/
mod.rs

1use crate::packed_rtree::{calc_extent, hilbert_sort, NodeItem, PackedRTree};
2use crate::MAGIC_BYTES;
3use attr_index::build_attribute_index_for_attr;
4use attribute::AttributeSchema;
5use cjseq::{CityJSON, CityJSONFeature, Transform as CjTransform};
6use feature_writer::{AttributeFeatureOffset, FeatureWriter};
7use header_writer::{HeaderWriter, HeaderWriterOptions};
8use serializer::AttributeIndexInfo;
9
10use crate::error::Result;
11use std::collections::BTreeMap;
12use std::fs::File;
13use std::io::{BufReader, BufWriter, Read, Seek, SeekFrom, Write};
14mod attr_index;
15pub mod attribute;
16pub mod error;
17pub mod feature_writer;
18pub mod geom_encoder;
19pub mod header_writer;
20pub mod serializer;
21/// Main writer for FlatCityBuf (FCB) format
22///
23/// FcbWriter handles the serialization of CityJSON data into the FCB binary format.
24/// It manages both header and feature writing, using a temporary file for feature storage
25/// before final assembly.
26pub struct FcbWriter<'a> {
27    /// Temporary buffer for storing features before final assembly
28    tmpout: BufWriter<File>,
29    /// Writer for the FCB header section
30    header_writer: HeaderWriter<'a>,
31    /// Optional writer for features
32    feat_writer: Option<FeatureWriter<'a>>,
33
34    transform: CjTransform,
35    /// Offset of the feature in the feature data section
36    feat_offsets: Vec<FeatureOffset>,
37    feat_nodes: Vec<NodeItem>,
38    attr_schema: AttributeSchema,
39    semantic_attr_schema: Option<AttributeSchema>,
40    // Temporary storage for attribute index entries. `BTreeMap`, not
41    // `HashMap`: `build_attribute_index_for_attr` walks it, and for keys that
42    // repeat across features the walk order becomes the payload-list order in
43    // the serialised B+tree. A randomly seeded order there is a different file
44    // on every run.
45    attribute_index_entries: BTreeMap<usize, AttributeFeatureOffset>,
46}
47
48#[derive(Clone, PartialEq, Debug)]
49struct FeatureOffset {
50    temp_feature_id: usize,
51    offset: usize,
52    size: usize,
53}
54
55impl<'a> FcbWriter<'a> {
56    /// Creates a new FCB writer instance
57    ///
58    /// # Arguments
59    ///
60    /// * `cj` - The CityJSON data to be written
61    /// * `header_option` - Optional configuration for header writing
62    /// * `first_feature` - Optional first feature to begin writing
63    ///
64    /// # Returns
65    ///
66    /// A Result containing the new FcbWriter instance
67    pub fn new(
68        cj: CityJSON,
69        header_option: Option<HeaderWriterOptions>,
70        attr_schema: Option<AttributeSchema>,
71        semantic_attr_schema: Option<AttributeSchema>,
72    ) -> Result<Self> {
73        let attr_schema = attr_schema.unwrap_or_default();
74
75        let transform = cj.transform.clone();
76        let header_writer = HeaderWriter::new(
77            cj,
78            header_option,
79            attr_schema.clone(),
80            semantic_attr_schema.clone(),
81        );
82        Ok(Self {
83            header_writer,
84            transform,
85            feat_writer: None,
86            tmpout: BufWriter::new(tempfile::tempfile()?),
87            attr_schema,
88            semantic_attr_schema,
89            feat_offsets: Vec::new(),
90            feat_nodes: Vec::new(),
91            attribute_index_entries: BTreeMap::new(),
92        })
93    }
94
95    /// Writes the current feature to the temporary buffer
96    ///
97    /// # Returns
98    ///
99    /// A Result indicating success or failure of the write operation
100    fn write_feature(&mut self) -> Result<()> {
101        let transform = &self.transform;
102
103        if let Some(feat_writer) = &mut self.feat_writer {
104            let feat_buf = feat_writer.finish_to_feature();
105
106            let mut attr_feature_offset = feat_writer.attribute_feature_offsets.clone();
107
108            let mut node = Self::actual_bbox(transform, &feat_writer.bbox);
109            node.offset = self.feat_offsets.len() as u64;
110            self.feat_nodes.push(node);
111
112            let tempoffset = self
113                .feat_offsets
114                .last()
115                .map(|it| it.offset + it.size)
116                .unwrap_or(0);
117
118            attr_feature_offset.offset = tempoffset;
119            self.attribute_index_entries
120                .insert(self.feat_offsets.len(), attr_feature_offset);
121
122            self.feat_offsets.push(FeatureOffset {
123                temp_feature_id: self.feat_offsets.len(),
124                offset: tempoffset,
125                size: feat_buf.len(),
126            });
127
128            self.tmpout.write_all(&feat_buf)?;
129        }
130        Ok(())
131    }
132
133    fn actual_bbox(transform: &CjTransform, bbox: &NodeItem) -> NodeItem {
134        let scale_x = transform.scale[0];
135        let scale_y = transform.scale[1];
136        let translate_x = transform.translate[0];
137        let translate_y = transform.translate[1];
138        NodeItem::bounds(
139            bbox.min_x * scale_x + translate_x,
140            bbox.min_y * scale_y + translate_y,
141            bbox.max_x * scale_x + translate_x,
142            bbox.max_y * scale_y + translate_y,
143        )
144    }
145
146    /// Adds a new feature to be written
147    ///
148    /// # Arguments
149    ///
150    /// * `feature` - The CityJSON feature to add
151    ///
152    /// # Returns
153    ///
154    /// A Result indicating success or failure of the operation
155    pub fn add_feature(&mut self, feature: &'a CityJSONFeature) -> Result<()> {
156        if self.feat_writer.is_none() {
157            self.feat_writer = Some(FeatureWriter::new(
158                feature,
159                self.attr_schema.clone(),
160                self.semantic_attr_schema.clone(),
161                self.header_writer
162                    .header_options
163                    .attribute_indices
164                    .as_ref()
165                    .map(|a| a.iter().map(|(name, _)| name.clone()).collect()),
166            ));
167        }
168
169        if let Some(feat_writer) = &mut self.feat_writer {
170            feat_writer.add_feature(feature);
171            self.write_feature()?;
172        }
173
174        Ok(())
175    }
176
177    /// Writes the complete FCB dataset to the output
178    ///
179    /// This method assembles the final FCB file by writing:
180    /// 1. Magic bytes
181    /// 2. Header
182    /// 3. Feature data
183    ///
184    /// # Arguments
185    ///
186    /// * `out` - The output destination implementing Write
187    ///
188    /// # Returns
189    ///
190    /// A Result indicating success or failure of the write operation
191    pub fn write(mut self, mut out: impl Write) -> Result<()> {
192        let mut attr_indices = self.header_writer.header_options.attribute_indices.clone();
193
194        // sort attribute indices by schema index (ascending)
195        if let Some(ref mut indices) = attr_indices {
196            indices.sort_by_key(|(name, _)| {
197                self.attr_schema
198                    .get(name)
199                    .map(|(idx, _)| *idx)
200                    .unwrap_or(u16::MAX)
201            });
202        }
203
204        out.write_all(&MAGIC_BYTES)?;
205        let index_node_size = self.header_writer.header_options.index_node_size;
206
207        let mut rtree_buf = Vec::new();
208        if index_node_size > 0 && !self.feat_nodes.is_empty() {
209            let extent = calc_extent(&self.feat_nodes);
210            hilbert_sort(&mut self.feat_nodes, &extent);
211            let mut offset = 0;
212            let index_nodes = self
213                .feat_nodes
214                .iter()
215                .map(|temp_node| {
216                    let feat = &self.feat_offsets[temp_node.offset as usize];
217                    let mut node = temp_node.clone();
218                    node.offset = offset;
219                    offset += feat.size as u64;
220                    node
221                })
222                .collect::<Vec<_>>();
223            let tree = PackedRTree::build(&index_nodes, &extent, index_node_size)?;
224            tree.stream_write(&mut rtree_buf)?;
225        }
226
227        self.tmpout.rewind()?;
228        let unsorted_feature_output = self.tmpout.into_inner().map_err(|e| e.into_error())?;
229        let mut unsorted_feature_reader = BufReader::new(unsorted_feature_output);
230
231        let mut sorted_feature_buf = Vec::with_capacity(2048);
232
233        for node in &self.feat_nodes {
234            let feat = &self.feat_offsets[node.offset as usize];
235            unsorted_feature_reader.seek(SeekFrom::Start(feat.offset as u64))?;
236
237            if let Some(attr_index_entry) =
238                self.attribute_index_entries.get_mut(&feat.temp_feature_id)
239            {
240                attr_index_entry.offset = sorted_feature_buf.len();
241                attr_index_entry.size = feat.size;
242            }
243
244            let cur_len = sorted_feature_buf.len();
245            sorted_feature_buf.resize(cur_len + feat.size, 0);
246            unsorted_feature_reader.read_exact(&mut sorted_feature_buf[cur_len..])?;
247        }
248
249        // build attribute index buffers in sorted order
250        let mut attr_index_buf: Vec<u8> = Vec::new();
251        let mut attr_index_info: Vec<AttributeIndexInfo> = Vec::new();
252        if let Some(sorted_indices) = &attr_indices {
253            for (name, bf_opt) in sorted_indices {
254                let bf = bf_opt.unwrap_or(crate::static_btree::DEFAULT_BRANCHING_FACTOR);
255                if let Ok((buf, info)) = build_attribute_index_for_attr(
256                    name,
257                    &self.attr_schema,
258                    &self.attribute_index_entries,
259                    bf,
260                ) {
261                    attr_index_info.push(info);
262                    attr_index_buf.extend(&buf);
263                }
264            }
265        }
266
267        // write header with attribute indices metadata
268        self.header_writer.attribute_indices_info = Some(attr_index_info);
269        let header_buf = self.header_writer.finish_to_header()?;
270        out.write_all(&header_buf)?;
271
272        // write spatial index (if any), attribute index bytes, then feature data
273        out.write_all(&rtree_buf)?;
274        out.write_all(&attr_index_buf)?;
275        out.write_all(&sorted_feature_buf)?;
276
277        Ok(())
278    }
279}