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;
21pub struct FcbWriter<'a> {
27 tmpout: BufWriter<File>,
29 header_writer: HeaderWriter<'a>,
31 feat_writer: Option<FeatureWriter<'a>>,
33
34 transform: CjTransform,
35 feat_offsets: Vec<FeatureOffset>,
37 feat_nodes: Vec<NodeItem>,
38 attr_schema: AttributeSchema,
39 semantic_attr_schema: Option<AttributeSchema>,
40 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 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 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 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 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 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 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 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 out.write_all(&rtree_buf)?;
274 out.write_all(&attr_index_buf)?;
275 out.write_all(&sorted_feature_buf)?;
276
277 Ok(())
278 }
279}