From 63bbbcdea3299679056c27ee90d310032ee6f1d0 Mon Sep 17 00:00:00 2001
From: Nick Lockwood <nick@charcoaldesign.co.uk>
Date: Sun, 20 Aug 2023 22:48:20 +0100
Subject: [PATCH] Add Mesh.stlData() export options

---
 Sources/Mesh+STL.swift      | 56 +++++++++++++++++++++++++++++++++----
 Tests/MeshExportTests.swift | 24 ++++++++++++++++
 2 files changed, 75 insertions(+), 5 deletions(-)

diff --git a/Sources/Mesh+STL.swift b/Sources/Mesh+STL.swift
index 05067483..6b4fade0 100644
--- a/Sources/Mesh+STL.swift
+++ b/Sources/Mesh+STL.swift
@@ -33,6 +33,35 @@ public struct STLTextOptions {
     }
 }
 
+/// Configuration options for binary STL export.
+public struct STLBinaryOptions {
+    /// Data to use for file header.
+    /// Note: data will be padded to 80 bytes. If more than 80 bytes are provided, data will be truncated.
+    public var header: Data
+    /// Should normal values be zeroed out?
+    public var zeroNormals: Bool
+    /// A closure that maps each polygon's material to an STL facet color.
+    public var colorLookup: Mesh.STLColorProvider
+
+    public init(
+        header: Data = .init(),
+        zeroNormals: Bool = false,
+        colorLookup: Mesh.STLColorProvider? = nil
+    ) {
+        self.header = header
+        self.zeroNormals = zeroNormals
+        self.colorLookup = colorLookup ?? defaultColorMapping
+    }
+}
+
+/// Configuration for exported STL file.
+public enum STLFormat {
+    /// Export in ASCII format.
+    case text(STLTextOptions)
+    /// Export in binary format.
+    case binary(STLBinaryOptions)
+}
+
 public extension Mesh {
     /// Return ASCII STL string data for the mesh.
     func stlString(name: String = "") -> String {
@@ -59,15 +88,32 @@ public extension Mesh {
     /// - Parameter colorLookup: A closure to map Euclid materials to STL facet colors. Use `nil` for default mapping.
     /// - Returns: A Euclid `Color` value.
     func stlData(colorLookup: STLColorProvider? = nil) -> Data {
+        stlData(options: .init(colorLookup: colorLookup))
+    }
+
+    /// Return binary STL data for the mesh.
+    /// - Parameter options: The output ooptions for the STL file
+    /// - Returns: The encoded STL data.
+    func stlData(options: STLBinaryOptions) -> Data {
         let triangles = triangulate().polygons
         let bufferSize = headerSize + 4 + triangles.count * triangleSize
         let buffer = Buffer(capacity: bufferSize)
+        options.header.copyBytes(to: buffer.buffer, count: min(options.header.count, 80))
         buffer.count = headerSize
         buffer.append(UInt32(triangles.count))
-        let colorLookup = colorLookup ?? defaultColorMapping
-        triangles.forEach { buffer.append($0, colorLookup: colorLookup) }
+        triangles.forEach { buffer.append($0, options: options) }
         return Data(buffer)
     }
+
+    /// Return STL data for the mesh.
+    func stlData(format: STLFormat) -> Data {
+        switch format {
+        case let .binary(options):
+            return stlData(options: options)
+        case let .text(options):
+            return Data(stlString(options: options).utf8)
+        }
+    }
 }
 
 private extension Polygon {
@@ -111,10 +157,10 @@ private extension Buffer {
         append(0x8000 | red << 10 | green << 5 | blue)
     }
 
-    func append(_ polygon: Polygon, colorLookup: Mesh.STLColorProvider) {
-        append(polygon.plane.normal)
+    func append(_ polygon: Polygon, options: STLBinaryOptions) {
+        append(options.zeroNormals ? .zero : polygon.plane.normal)
         polygon.vertices.forEach { append($0.position) }
-        if let color = colorLookup(polygon.material) {
+        if let color = options.colorLookup(polygon.material) {
             append(color)
         } else {
             count += 2
diff --git a/Tests/MeshExportTests.swift b/Tests/MeshExportTests.swift
index 301bf3d0..e25a342a 100644
--- a/Tests/MeshExportTests.swift
+++ b/Tests/MeshExportTests.swift
@@ -149,6 +149,30 @@ class MeshExportTests: XCTestCase {
         """)
     }
 
+    func testCubeSTLDataWithCustomHeader() {
+        let cube = Mesh.cube().translated(by: Vector(0.5, 0.5, 0.5))
+        let header = "Hello World".data(using: .utf8)!
+        let stlData = cube.stlData(options: .init(header: header))
+        XCTAssertEqual(stlData.count, 80 + 4 + 12 * 50)
+        XCTAssertEqual(stlData.prefix(header.count), header)
+        let hex = stlData.reduce(into: "") { $0 += String(format: "%02x", $1) }
+        XCTAssertEqual(hex, """
+        48656c6c6f20576f726c6400000000000000000000000000000000000000000000000000000000000000000000000000000000000000000\
+        00000000000000000000000000000000000000000000000000c0000000000803f00000000000000000000803f000000000000803f000080\
+        3f00000000000000000000803f0000803f0000000000000000803f00000000000000000000803f000000000000803f0000803f0000803f0\
+        00000000000803f0000803f0000803f0000000080bf000000800000008000000000000000000000000000000000000000000000803f0000\
+        00000000803f0000803f0000000080bf0000008000000080000000000000000000000000000000000000803f0000803f000000000000803\
+        f000000000000000000000000803f00000000000000000000803f0000803f0000803f0000803f0000803f0000803f0000803f0000000000\
+        00000000000000803f00000000000000000000803f0000803f0000803f0000803f00000000000000000000803f000000000000000000800\
+        00080bf000000800000000000000000000000000000803f00000000000000000000803f000000000000803f000000000080000080bf0000\
+        00800000000000000000000000000000803f000000000000803f00000000000000000000803f000000000000000000000000803f0000000\
+        0000000000000803f0000803f000000000000803f0000803f0000803f0000803f000000000000000000000000803f000000000000000000\
+        00803f0000803f0000803f0000803f000000000000803f0000803f00000000008000000080000080bf0000803f000000000000000000000\
+        0000000000000000000000000000000803f0000000000000000008000000080000080bf0000803f0000000000000000000000000000803f\
+        000000000000803f0000803f000000000000
+        """)
+    }
+
     // MARK: OBJ export
 
     func testCubeOBJ() {