main.go

package main

import (
	"fmt"
	"math"
	"os"
)

const (
	imageWidth  = 1024
	imageHeight = 768
	fov         = 90 * math.Pi / 180
	maxDepth    = 4
	epsilon     = float32(1e-3)
)

type Vec3 struct {
	x, y, z float32
}

func (v Vec3) Add(u Vec3) Vec3 {
	return Vec3{v.x + u.x, v.y + u.y, v.z + u.z}
}

func (v Vec3) Sub(u Vec3) Vec3 {
	return Vec3{v.x - u.x, v.y - u.y, v.z - u.z}
}

func (v Vec3) Mul(a float32) Vec3 {
	return Vec3{v.x * a, v.y * a, v.z * a}
}

func (v Vec3) Div(a float32) Vec3 {
	return Vec3{v.x / a, v.y / a, v.z / a}
}

func (v Vec3) Dot(u Vec3) float32 {
	return v.x*u.x + v.y*u.y + v.z*u.z
}

func (v Vec3) LengthSquared() float32 {
	return v.Dot(v)
}

func (v Vec3) Normalize() Vec3 {
	length := float32(math.Sqrt(float64(v.LengthSquared())))
	if length == 0 {
		return v
	}
	return Vec3{v.x / length, v.y / length, v.z / length}
}

func clamp01(v float32) float32 {
	if v < 0 {
		return 0
	}
	if v > 1 {
		return 1
	}
	return v
}

func max3(a, b, c float32) float32 {
	max := a
	if b > max {
		max = b
	}
	if c > max {
		max = c
	}
	return max
}

type Material struct {
	diffuseColor     Vec3
	specularExponent float32
	albedo           Vec3
}

type Sphere struct {
	center   Vec3
	radius   float32
	material Material
}

type Light struct {
	position  Vec3
	intensity float32
}

func reflectVec(inLight, surfaceNormal Vec3) Vec3 {
	dot := inLight.Dot(surfaceNormal)
	return inLight.Sub(surfaceNormal.Mul(2 * dot))
}

func (s Sphere) intersectRay(orig Vec3, dir Vec3) (bool, float32) {
	L := s.center.Sub(orig)
	tca := L.Dot(dir)
	d2 := L.Dot(L) - tca*tca
	if d2 > s.radius*s.radius {
		return false, 0
	}
	thc := float32(math.Sqrt(float64(s.radius*s.radius - d2)))
	t0 := tca - thc
	t1 := tca + thc
	if t0 < 0 {
		t0 = t1
	}
	if t0 < 0 {
		return false, 0
	}
	return true, t0
}

func sceneIntersect(orig Vec3, dir Vec3, spheres []Sphere) (bool, Vec3, Vec3, Material) {
	sphereDist := float32(math.MaxFloat32)
	var hitPoint, hitNormal Vec3
	var hitMaterial Material
	hit := false

	for i := range spheres {
		sphere := spheres[i]
		isIntersect, dist := sphere.intersectRay(orig, dir)
		if isIntersect && dist < sphereDist {
			sphereDist = dist
			hitPoint = orig.Add(dir.Mul(dist))
			hitNormal = hitPoint.Sub(sphere.center).Normalize()
			hitMaterial = sphere.material
			hit = true
		}
	}
	return hit, hitPoint, hitNormal, hitMaterial
}

func castRay(orig Vec3, dir Vec3, spheres []Sphere, lights []Light, depth int) Vec3 {
	hit, hitPoint, hitNormal, material := sceneIntersect(orig, dir, spheres)
	if depth > maxDepth || !hit {
		return Vec3{0.2, 0.7, 0.8}
	}

	reflectDir := reflectVec(dir, hitNormal).Normalize()
	var reflectOrig Vec3
	if reflectDir.Dot(hitNormal) < 0 {
		reflectOrig = hitPoint.Sub(hitNormal.Mul(epsilon))
	} else {
		reflectOrig = hitPoint.Add(hitNormal.Mul(epsilon))
	}
	reflectColor := castRay(reflectOrig, reflectDir, spheres, lights, depth+1)

	var diffuseLightIntensity float32 = 0
	var specularLightIntensity float32 = 0
	for i := range lights {
		light := lights[i]
		lightVec := light.position.Sub(hitPoint)
		lightDir := lightVec.Normalize()
		lightDist := lightVec.LengthSquared()

		var shadowOrig Vec3
		if lightDir.Dot(hitNormal) < 0 {
			shadowOrig = hitPoint.Sub(hitNormal.Mul(epsilon))
		} else {
			shadowOrig = hitPoint.Add(hitNormal.Mul(epsilon))
		}
		shadowHit, shadowPoint, _, _ := sceneIntersect(shadowOrig, lightDir, spheres)
		if shadowHit && shadowPoint.Sub(shadowOrig).LengthSquared() < lightDist {
			continue
		}
		diffuseLightIntensity += light.intensity * float32(math.Max(0.0, float64(lightDir.Dot(hitNormal))))
		reflectDir := reflectVec(lightDir.Mul(-1), hitNormal)
		viewDir := dir.Mul(-1)
		dot := reflectDir.Dot(viewDir)
		if dot > 0 {
			specularLightIntensity += light.intensity * float32(math.Pow(float64(dot), float64(material.specularExponent)))
		}
	}
	diffuse := material.diffuseColor.Mul(diffuseLightIntensity).Mul(material.albedo.x)
	specular := Vec3{1, 1, 1}.Mul(specularLightIntensity).Mul(material.albedo.y)
	reflection := reflectColor.Mul(material.albedo.z)
	return diffuse.Add(specular).Add(reflection)
}

func render(spheres []Sphere, lights []Light) error {
	frameBuffer := make([]Vec3, imageHeight*imageWidth)

	camera := Vec3{0, 0, 0}
	aspect := float32(imageWidth) / float32(imageHeight)
	tanFov := float32(math.Tan(fov / 2))

	for j := 0; j < imageHeight; j++ {
		for i := 0; i < imageWidth; i++ {
			x := (2*(float32(i)+0.5)/float32(imageWidth) - 1) * tanFov * aspect
			y := -(2*(float32(j)+0.5)/float32(imageHeight) - 1) * tanFov
			dir := Vec3{x, y, -1}.Normalize()
			frameBuffer[i+j*imageWidth] = castRay(camera, dir, spheres, lights, 0)
		}
	}

	header := fmt.Sprintf("P6\n%d %d\n255\n", imageWidth, imageHeight)
	dataSize := imageWidth * imageHeight * 3

	buffer := make([]byte, 0, len(header)+dataSize)
	buffer = append(buffer, header...)

	for i := 0; i < imageWidth*imageHeight; i++ {
		color := frameBuffer[i]
		maxValue := max3(color.x, color.y, color.z)
		if maxValue > 1 {
			color = color.Div(maxValue)
		}
		buffer = append(buffer, byte(clamp01(color.x)*255))
		buffer = append(buffer, byte(clamp01(color.y)*255))
		buffer = append(buffer, byte(clamp01(color.z)*255))
	}

	return os.WriteFile("./out.ppm", buffer, 0644)
}

func main() {
	fmt.Println("tiny ray tracing in golang")

	ivory := Material{
		diffuseColor:     Vec3{0.4, 0.4, 0.3},
		specularExponent: 50,
		albedo:           Vec3{0.6, 0.3, 0.1},
	}
	redRubber := Material{
		diffuseColor:     Vec3{0.3, 0.1, 0.1},
		specularExponent: 10,
		albedo:           Vec3{0.9, 0.1, 0.0},
	}
	mirror := Material{
		diffuseColor:     Vec3{1.0, 1.0, 1.0},
		specularExponent: 10,
		albedo:           Vec3{0.0, 10.0, 0.8},
	}

	spheres := []Sphere{
		{center: Vec3{-3, 0, -16}, radius: 2, material: ivory},
		{center: Vec3{-1.0, -1.5, -12}, radius: 2, material: mirror},
		{center: Vec3{1.5, -0.5, -18}, radius: 3, material: redRubber},
		{center: Vec3{7, 5, -18}, radius: 4, material: mirror},
	}

	lights := []Light{
		{position: Vec3{-20, 20, 20}, intensity: 1.5},
		{position: Vec3{30, 50, -25}, intensity: 1.8},
		{position: Vec3{30, 20, 30}, intensity: 1.7},
	}

	if err := render(spheres, lights); err != nil {
		fmt.Fprintln(os.Stderr, "render:", err)
		os.Exit(1)
	}
}