Skip to main content

draco_io/
stl_writer.rs

1//! STL format writer for triangle meshes.
2//!
3//! Writes both containers: binary, which is the default because it is a fifth
4//! the size and exact, and ASCII for the cases that call for readable output.
5//!
6//! # Example
7//!
8//! ```no_run
9//! use draco_io::{StlWriter, Writer};
10//!
11//! let mesh = draco_core::mesh::Mesh::new();
12//! let mut writer = StlWriter::new();
13//! writer.add_mesh(&mesh, Some("Model"))?;
14//! writer.write("output.stl")?;
15//! # Ok::<(), std::io::Error>(())
16//! ```
17
18use 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/// Which STL container to write.
29#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
30pub enum StlFormat {
31    /// The 80-byte header and 50-byte triangle records.
32    #[default]
33    Binary,
34    /// `solid`/`facet`/`vertex` text.
35    Ascii,
36}
37
38/// STL format writer.
39#[derive(Debug, Clone, Default)]
40pub struct StlWriter {
41    format: StlFormat,
42    /// The name written into the ASCII `solid` line; binary has no field for it.
43    name: String,
44    triangles: Vec<[[f32; 3]; 3]>,
45}
46
47impl StlWriter {
48    /// Choose the container to write. Binary is the default.
49    pub fn with_format(mut self, format: StlFormat) -> Self {
50        self.format = format;
51        self
52    }
53
54    /// Number of triangles collected so far.
55    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    /// Add a mesh's triangles.
66    ///
67    /// STL holds geometry and nothing else — no vertex identity, no attributes,
68    /// one facet normal per triangle — so only positions and the face list are
69    /// read. Several meshes may be added; they become one solid, which is all
70    /// the format can represent.
71    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        // The header is free-form, and a reader deciding the container by its
131        // first word would read this file as ASCII if it began with `solid`.
132        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
173/// The facet normal, from the winding rather than from any stored normal.
174///
175/// STL states one normal per facet and readers take it as the surface's own, so
176/// a smooth-shading normal averaged from the corners would contradict the
177/// geometry it is attached to. A degenerate triangle has no normal to state, and
178/// zero is the format's way of saying the reader should derive one.
179fn 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        // Counter-clockwise in the XY plane, so the facet normal is +Z.
251        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        // And the header must not start a binary file with the ASCII keyword.
259        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    /// Both containers have to come back as the same geometry, and the binary
275    /// one exactly: it stores the float bits the mesh held.
276    #[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}