agent-tools/mesh-tool/mesh/stl.go

package mesh

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"math"
	"strconv"
	"strings"
)

// ReadSTL reads binary or ASCII STL (auto-detected). STL stores loose
// triangles, so identical vertices are welded back together to recover
// connectivity (needed for watertight checks and sane OBJ export).
func ReadSTL(r io.Reader) (*Scene, error) {
	data, err := io.ReadAll(r)
	if err != nil {
		return nil, err
	}
	if len(data) >= 84 {
		n := binary.LittleEndian.Uint32(data[80:84])
		if int(84+50*n) == len(data) {
			return readSTLBinary(data)
		}
	}
	if bytes.HasPrefix(bytes.TrimLeft(data, " \t\r\n"), []byte("solid")) {
		return readSTLAscii(data)
	}
	return nil, fmt.Errorf("not a valid STL file (neither binary layout nor 'solid ...' text)")
}

type welder struct {
	mesh  *Mesh
	index map[Vec3]int
}

func newWelder(name string) *welder {
	return &welder{mesh: &Mesh{Name: name}, index: map[Vec3]int{}}
}

func (w *welder) add(a, b, c Vec3) {
	var t Triangle
	for i, v := range [3]Vec3{a, b, c} {
		idx, ok := w.index[v]
		if !ok {
			idx = len(w.mesh.Verts)
			w.mesh.Verts = append(w.mesh.Verts, v)
			w.index[v] = idx
		}
		t[i] = idx
	}
	if t[0] == t[1] || t[1] == t[2] || t[2] == t[0] {
		return // degenerate
	}
	w.mesh.Tris = append(w.mesh.Tris, t)
}

func readSTLBinary(data []byte) (*Scene, error) {
	n := int(binary.LittleEndian.Uint32(data[80:84]))
	w := newWelder("stl")
	off := 84
	f32 := func(o int) float64 {
		return float64(math.Float32frombits(binary.LittleEndian.Uint32(data[o : o+4])))
	}
	for i := 0; i < n; i++ {
		// 12 bytes normal (ignored), 3 * 12 bytes vertices, 2 bytes attrs
		var v [3]Vec3
		for k := 0; k < 3; k++ {
			base := off + 12 + k*12
			v[k] = Vec3{f32(base), f32(base + 4), f32(base + 8)}
		}
		w.add(v[0], v[1], v[2])
		off += 50
	}
	return &Scene{Meshes: []*Mesh{w.mesh}}, nil
}

func readSTLAscii(data []byte) (*Scene, error) {
	name := "stl"
	w := newWelder(name)
	var cur []Vec3
	sc := bufio.NewScanner(bytes.NewReader(data))
	sc.Buffer(make([]byte, 0, 64*1024), 16*1024*1024)
	lineNo := 0
	for sc.Scan() {
		lineNo++
		fields := strings.Fields(sc.Text())
		if len(fields) == 0 {
			continue
		}
		switch fields[0] {
		case "solid":
			if len(fields) > 1 {
				w.mesh.Name = fields[1]
			}
		case "vertex":
			if len(fields) < 4 {
				return nil, fmt.Errorf("stl line %d: vertex needs x y z", lineNo)
			}
			var v Vec3
			var err error
			if v.X, err = strconv.ParseFloat(fields[1], 64); err == nil {
				if v.Y, err = strconv.ParseFloat(fields[2], 64); err == nil {
					v.Z, err = strconv.ParseFloat(fields[3], 64)
				}
			}
			if err != nil {
				return nil, fmt.Errorf("stl line %d: bad vertex: %v", lineNo, err)
			}
			cur = append(cur, v)
		case "endfacet":
			if len(cur) != 3 {
				return nil, fmt.Errorf("stl line %d: facet has %d vertices, want 3", lineNo, len(cur))
			}
			w.add(cur[0], cur[1], cur[2])
			cur = cur[:0]
		}
	}
	if err := sc.Err(); err != nil {
		return nil, err
	}
	if len(w.mesh.Tris) == 0 {
		return nil, fmt.Errorf("stl contains no triangles")
	}
	return &Scene{Meshes: []*Mesh{w.mesh}}, nil
}

// WriteSTLBinary writes the whole scene as one binary STL solid
// (STL has no concept of multiple named objects).
func WriteSTLBinary(w io.Writer, s *Scene) error {
	m := s.Merged()
	bw := bufio.NewWriter(w)
	header := make([]byte, 80)
	copy(header, []byte("exported by mesht (agent-tools)"))
	bw.Write(header)
	binary.Write(bw, binary.LittleEndian, uint32(len(m.Tris)))
	buf := make([]byte, 50)
	for i, t := range m.Tris {
		n := m.FaceNormal(i)
		le := binary.LittleEndian
		le.PutUint32(buf[0:], math.Float32bits(float32(n.X)))
		le.PutUint32(buf[4:], math.Float32bits(float32(n.Y)))
		le.PutUint32(buf[8:], math.Float32bits(float32(n.Z)))
		for k := 0; k < 3; k++ {
			v := m.Verts[t[k]]
			le.PutUint32(buf[12+k*12:], math.Float32bits(float32(v.X)))
			le.PutUint32(buf[16+k*12:], math.Float32bits(float32(v.Y)))
			le.PutUint32(buf[20+k*12:], math.Float32bits(float32(v.Z)))
		}
		buf[48], buf[49] = 0, 0
		bw.Write(buf)
	}
	return bw.Flush()
}

// WriteSTLAscii writes the scene as a text STL solid.
func WriteSTLAscii(w io.Writer, s *Scene) error {
	m := s.Merged()
	bw := bufio.NewWriter(w)
	fmt.Fprintf(bw, "solid %s\n", sanitizeToken(m.Name))
	for i, t := range m.Tris {
		n := m.FaceNormal(i)
		fmt.Fprintf(bw, "  facet normal %g %g %g\n    outer loop\n", n.X, n.Y, n.Z)
		for k := 0; k < 3; k++ {
			v := m.Verts[t[k]]
			fmt.Fprintf(bw, "      vertex %g %g %g\n", v.X, v.Y, v.Z)
		}
		fmt.Fprintf(bw, "    endloop\n  endfacet\n")
	}
	fmt.Fprintf(bw, "endsolid %s\n", sanitizeToken(m.Name))
	return bw.Flush()
}

func sanitizeToken(s string) string {
	s = strings.ReplaceAll(strings.TrimSpace(s), " ", "_")
	if s == "" {
		return "mesh"
	}
	return s
}