1use std::fs::File;
19use std::io::{self, BufWriter, Write};
20use std::path::Path;
21
22use draco_core::geometry_attribute::GeometryAttributeType;
23use draco_core::geometry_indices::FaceIndex;
24use draco_core::mesh::Mesh;
25
26use crate::traits::{WriteToBytes, Writer};
27
28#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
30pub enum StlFormat {
31 #[default]
33 Binary,
34 Ascii,
36}
37
38#[derive(Debug, Clone, Default)]
40pub struct StlWriter {
41 format: StlFormat,
42 name: String,
44 triangles: Vec<[[f32; 3]; 3]>,
45}
46
47impl StlWriter {
48 pub fn with_format(mut self, format: StlFormat) -> Self {
50 self.format = format;
51 self
52 }
53
54 pub fn triangle_count(&self) -> usize {
56 self.triangles.len()
57 }
58}
59
60impl Writer for StlWriter {
61 fn new() -> Self {
62 Self::default()
63 }
64
65 fn add_mesh(&mut self, mesh: &Mesh, name: Option<&str>) -> io::Result<()> {
72 if self.name.is_empty() {
73 if let Some(name) = name {
74 self.name = name.to_string();
75 }
76 }
77 let position_id = mesh.named_attribute_id(GeometryAttributeType::Position);
78 if position_id < 0 {
79 return Err(io::Error::new(
80 io::ErrorKind::InvalidData,
81 "Mesh has no position attribute",
82 ));
83 }
84 let point_count = mesh.num_points();
85 for index in 0..mesh.num_faces() as u32 {
86 let face = mesh.face(FaceIndex(index));
87 if face.iter().any(|point| point.0 as usize >= point_count) {
88 return Err(io::Error::new(
89 io::ErrorKind::InvalidData,
90 "Mesh face references a point outside the attribute",
91 ));
92 }
93 self.triangles.push([
94 read_position(mesh, position_id, face[0].0 as usize),
95 read_position(mesh, position_id, face[1].0 as usize),
96 read_position(mesh, position_id, face[2].0 as usize),
97 ]);
98 }
99 Ok(())
100 }
101
102 fn write<P: AsRef<Path>>(&self, path: P) -> io::Result<()> {
103 let mut file = BufWriter::new(File::create(path)?);
104 file.write_all(&self.write_to_vec()?)?;
105 file.flush()
106 }
107
108 fn vertex_count(&self) -> usize {
109 self.triangles.len() * 3
110 }
111
112 fn face_count(&self) -> usize {
113 self.triangles.len()
114 }
115}
116
117impl WriteToBytes for StlWriter {
118 fn write_to_vec(&self) -> io::Result<Vec<u8>> {
119 match self.format {
120 StlFormat::Binary => Ok(self.write_binary()),
121 StlFormat::Ascii => Ok(self.write_ascii().into_bytes()),
122 }
123 }
124}
125
126impl StlWriter {
127 fn write_binary(&self) -> Vec<u8> {
128 let mut bytes = Vec::with_capacity(84 + self.triangles.len() * 50);
129 let mut header = [0u8; 80];
130 let label = format!("Draco {}", self.name);
133 let label = label.trim_end().as_bytes();
134 let length = label.len().min(80);
135 header[..length].copy_from_slice(&label[..length]);
136 bytes.extend_from_slice(&header);
137 bytes.extend_from_slice(&(self.triangles.len() as u32).to_le_bytes());
138 for triangle in &self.triangles {
139 for component in facet_normal(triangle) {
140 bytes.extend_from_slice(&component.to_le_bytes());
141 }
142 for vertex in triangle {
143 for component in vertex {
144 bytes.extend_from_slice(&component.to_le_bytes());
145 }
146 }
147 bytes.extend_from_slice(&0u16.to_le_bytes());
148 }
149 bytes
150 }
151
152 fn write_ascii(&self) -> String {
153 let name = if self.name.is_empty() {
154 "mesh"
155 } else {
156 &self.name
157 };
158 let mut text = String::new();
159 text.push_str(&format!("solid {name}\n"));
160 for triangle in &self.triangles {
161 let [nx, ny, nz] = facet_normal(triangle);
162 text.push_str(&format!(" facet normal {nx} {ny} {nz}\n outer loop\n"));
163 for [x, y, z] in triangle {
164 text.push_str(&format!(" vertex {x} {y} {z}\n"));
165 }
166 text.push_str(" endloop\n endfacet\n");
167 }
168 text.push_str(&format!("endsolid {name}\n"));
169 text
170 }
171}
172
173fn facet_normal(triangle: &[[f32; 3]; 3]) -> [f32; 3] {
180 let [a, b, c] = triangle;
181 let u = [b[0] - a[0], b[1] - a[1], b[2] - a[2]];
182 let v = [c[0] - a[0], c[1] - a[1], c[2] - a[2]];
183 let normal = [
184 u[1] * v[2] - u[2] * v[1],
185 u[2] * v[0] - u[0] * v[2],
186 u[0] * v[1] - u[1] * v[0],
187 ];
188 let length = (normal[0] * normal[0] + normal[1] * normal[1] + normal[2] * normal[2]).sqrt();
189 if length > 0.0 && length.is_finite() {
190 [normal[0] / length, normal[1] / length, normal[2] / length]
191 } else {
192 [0.0, 0.0, 0.0]
193 }
194}
195
196fn read_position(mesh: &Mesh, attribute_id: i32, point: usize) -> [f32; 3] {
197 let attribute = mesh.attribute(attribute_id);
198 let stride = attribute.byte_stride() as usize;
199 let mut bytes = [0u8; 12];
200 attribute.buffer().read(point * stride, &mut bytes);
201 [
202 f32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]),
203 f32::from_le_bytes([bytes[4], bytes[5], bytes[6], bytes[7]]),
204 f32::from_le_bytes([bytes[8], bytes[9], bytes[10], bytes[11]]),
205 ]
206}
207
208#[cfg(test)]
209mod tests {
210 use super::*;
211 use draco_core::draco_types::DataType;
212 use draco_core::geometry_attribute::PointAttribute;
213 use draco_core::geometry_indices::PointIndex;
214
215 fn triangle_mesh() -> Mesh {
216 let mut mesh = Mesh::new();
217 mesh.set_num_points(3);
218 mesh.set_num_faces(1);
219 let mut attribute = PointAttribute::new();
220 attribute.init(
221 GeometryAttributeType::Position,
222 3,
223 DataType::Float32,
224 false,
225 3,
226 );
227 for (index, vertex) in [[0.0f32, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]
228 .iter()
229 .enumerate()
230 {
231 let bytes: Vec<u8> = vertex
232 .iter()
233 .flat_map(|value| value.to_le_bytes())
234 .collect();
235 attribute.buffer_mut().write(index * 12, &bytes);
236 }
237 mesh.add_attribute(attribute);
238 mesh.set_face(FaceIndex(0), [PointIndex(0), PointIndex(1), PointIndex(2)]);
239 mesh
240 }
241
242 #[test]
243 fn test_write_binary_stl() {
244 let mut writer = StlWriter::new();
245 writer.add_mesh(&triangle_mesh(), Some("Tri")).unwrap();
246 let bytes = writer.write_to_vec().unwrap();
247
248 assert_eq!(bytes.len(), 84 + 50);
249 assert_eq!(u32::from_le_bytes(bytes[80..84].try_into().unwrap()), 1);
250 let normal: Vec<f32> = (0..3)
252 .map(|index| {
253 let start = 84 + index * 4;
254 f32::from_le_bytes(bytes[start..start + 4].try_into().unwrap())
255 })
256 .collect();
257 assert_eq!(normal, vec![0.0, 0.0, 1.0]);
258 assert!(!bytes.starts_with(b"solid"));
260 }
261
262 #[test]
263 fn test_write_ascii_stl() {
264 let mut writer = StlWriter::new().with_format(StlFormat::Ascii);
265 writer.add_mesh(&triangle_mesh(), Some("Tri")).unwrap();
266 let text = String::from_utf8(writer.write_to_vec().unwrap()).unwrap();
267
268 assert!(text.starts_with("solid Tri\n"));
269 assert!(text.contains("facet normal 0 0 1"));
270 assert_eq!(text.matches("vertex ").count(), 3);
271 assert!(text.trim_end().ends_with("endsolid Tri"));
272 }
273
274 #[cfg(feature = "stl-reader")]
277 #[test]
278 fn test_roundtrip_through_both_containers() {
279 use crate::stl_reader::StlReader;
280
281 for format in [StlFormat::Binary, StlFormat::Ascii] {
282 let mut writer = StlWriter::new().with_format(format);
283 writer.add_mesh(&triangle_mesh(), Some("Tri")).unwrap();
284 let mesh = StlReader::read_from_bytes(&writer.write_to_vec().unwrap()).unwrap();
285
286 assert_eq!(mesh.num_faces(), 1, "{format:?}");
287 assert_eq!(mesh.num_points(), 3, "{format:?}");
288 let position_id = mesh.named_attribute_id(GeometryAttributeType::Position);
289 assert_eq!(
290 read_position(&mesh, position_id, 1),
291 [1.0, 0.0, 0.0],
292 "{format:?}"
293 );
294 assert_eq!(
295 read_position(&mesh, position_id, 2),
296 [0.0, 1.0, 0.0],
297 "{format:?}"
298 );
299 }
300 }
301
302 #[test]
303 fn test_write_rejects_a_face_pointing_past_its_vertices() {
304 let mut mesh = triangle_mesh();
305 mesh.set_face(FaceIndex(0), [PointIndex(0), PointIndex(1), PointIndex(9)]);
306 let mut writer = StlWriter::new();
307 let error = writer.add_mesh(&mesh, None).unwrap_err();
308 assert_eq!(error.kind(), io::ErrorKind::InvalidData);
309 }
310}