Add write_vtk functions and test for vtk string creation

zigsdof/src/stl.zig

@@ -117,6 +117,75 @@ pub fn translateMesh(mesh: *Mesh, p: vec3.Vec3F64) void {
     return;
 }
 
+pub fn writeVtk(allocator: std.mem.Allocator, io: std.Io, mesh: *Mesh, file_path: []const u8) !void {
+    const out_string = try createVtkString(allocator, mesh);
+    defer allocator.free(out_string);
+
+    const cwd = std.Io.Dir.cwd();
+    const file = try cwd.openFile(io, file_path, .{ .mode = .read_write });
+    defer file.close(io);
+
+    const buf: [266]u8 = undefined;
+    const writer = file.writer(io, buf);
+    try writer.interface.writeAll(out_string);
+    try writer.interface.flush();
+    return;
+}
+
+// Create VTK file contents from a mesh.
+// NOTE: caller is responsible for cleaning up memory.
+fn createVtkString(allocator: std.mem.Allocator, mesh: *Mesh) ![]const u8 {
+    var output = try std.ArrayList(u8).initCapacity(
+        allocator,
+        6 * 1000 + mesh.vertices.items.len * 16 + mesh.triangle_vertices.items.len * 16,
+    );
+    const header =
+        \\# vtk DataFile Version 2.0
+        \\
+        \\ASCII
+        \\DATASET POLYDATA
+        \\
+    ;
+    try output.appendSlice(allocator, header);
+    var buf: [256]u8 = undefined;
+    const points_header = try std.fmt.bufPrint(
+        &buf,
+        "POINTS {} float\n",
+        .{mesh.vertices.items.len},
+    );
+    try output.appendSlice(allocator, points_header);
+
+    for (mesh.vertices.items) |vertex| {
+        const vtx = try std.fmt.bufPrint(&buf, "{} {} {}\n", .{
+            vertex.x(),
+            vertex.y(),
+            vertex.z(),
+        });
+        try output.appendSlice(allocator, vtx);
+    }
+
+    const polydata_header = try std.fmt.bufPrint(
+        &buf,
+        "POLYGONS {} {}\n",
+        .{
+            mesh.triangle_vertices.items.len,
+            mesh.triangle_vertices.items.len * 4,
+        },
+    );
+    try output.appendSlice(allocator, polydata_header);
+
+    for (mesh.triangle_vertices.items) |tri| {
+        const triangle = try std.fmt.bufPrint(&buf, "3 {} {} {}\n", .{
+            tri[0],
+            tri[1],
+            tri[2],
+        });
+        try output.appendSlice(allocator, triangle);
+    }
+
+    return output.toOwnedSlice(allocator);
+}
+
 test "Verify fixed point conversion works" {
     const vec = vec3.Vec3F64.init(1.0, 2.134, 1010.11);
     const fixed_point = vecToFixedPoint(vec, 2);
@@ -269,3 +338,29 @@ test "mesh is translated correctly" {
         1e-12,
     ));
 }
+
+test "vtk string is written correctly" {
+    const allocator = std.testing.allocator;
+
+    var mesh = try createDummyMesh(allocator);
+    defer mesh.deinit(allocator);
+
+    const vtk_string = try createVtkString(allocator, &mesh);
+    defer allocator.free(vtk_string);
+
+    const expected =
+        \\# vtk DataFile Version 2.0
+        \\
+        \\ASCII
+        \\DATASET POLYDATA
+        \\POINTS 3 float
+        \\0 0 0
+        \\0.5 0.8660254037844386 0
+        \\1 0 0
+        \\POLYGONS 1 4
+        \\3 0 1 2
+        \\
+    ;
+
+    try std.testing.expect(std.mem.eql(u8, vtk_string, expected));
+}