// Package fractal iteratively renders and magnifies fractals of the Mandelbrot or filled Julia set.
//
// Tested on a MacBook Air M1.
package main
import (
"context"
_ "embed"
"flag"
"fmt"
"io/fs"
"log/slog"
"os"
"os/signal"
"regexp"
"strconv"
"strings"
"sync"
"sync/atomic"
"syscall"
"time"
"unsafe"
"github.com/gogpu/gg"
_ "github.com/gogpu/gg/gpu"
"github.com/gogpu/gg/integration/ggcanvas"
"github.com/gogpu/gg/text"
"github.com/gogpu/gogpu"
"github.com/gogpu/gpucontext"
"github.com/gogpu/gputypes"
"github.com/gogpu/wgpu"
"github.com/jbunds/progress"
)
//go:embed common.wgsl
var commonShaderCode string
//go:embed mandelbrot.wgsl
var mandelbrotShaderCode string
//go:embed julia.wgsl
var juliaShaderCode string
const (
mainWidth = 800 // primary window logical width
mainHeight = 800 // primary window logical height
aboutWidth = 240 // about window logical width
aboutHeight = 164 // about window logical height
baseIterations = 500 // initial number of maximum iterations used to compute interior boundaries
paletteSize = 2000 // number of colors to pre-compute and pass to the GPU shader for fast lookup
viewportWidth = 3.0 // viewport width of the initial frame, i.e., the span of the complex plane rendered to the viewport
scaleFactor = 0.993 // multiplicative factor by which each rendered frame is iteratively magnified
linearGrowthFactor = 0.3 // multiplicative factor by which interior boundary calculation maximum iterations increase per each successive Mandelbrot fractal frame
quadraticGrowthFactor = 0.003 // multiplicative factor by which interior boundary calculation maximum iterations increase per the square of the Julia fractal frame count
maxPrecisionFrames = 2745 // empirically-determined limit for the number of frames to render before reaching precision limit
)
// state stores low-level state (color palette, viewport width, rendered frame count, and FPS, and uniforms).
type state struct {
frameCount int // tracks the number of frames rendered
paletteColors []uint32 // pre-computed color palette
viewportWidth, // tracks the width of the current frame's view of the complex plane
fps float64 // fps imprecisely tracks FPS rendered
xRealHi, xRealLo,
yImagHi, yImagLo float32 // target coordinates of the Mandelbrot set defined by ๐‘“(๐‘ง) = ๐‘งยฒ + ๐‘
cRealHi, cRealLo,
cImagHi, cImagLo float32 // complex constant term of the filled Julia set defined by ๐‘“(๐‘ง) = ๐‘งยฒ + ๐‘
}
// gpu stores all GPU resources required for rendering (device, buffers, compute pipeline).
type gpu struct {
device *wgpu.Device // logical GPU device
paletteTex *wgpu.Texture // pre-computed color palette 1D texture
uniformBuf *wgpu.Buffer // uniforms buffer
staticBindGroup *wgpu.BindGroup // uniforms and color palette buffer
bgLayout0 *wgpu.BindGroupLayout // uniforms and color palette layout
bgLayout1 *wgpu.BindGroupLayout // storage texture layout
pipeline *wgpu.ComputePipeline // GPU compute pipeline configuration
shaderCode string // WGSL shader code (specific to fractal kind: Mandelbrot or Julia)
}
// assets stores assets used to render frames (canvas, font, texture view).
type assets struct {
canvas *ggcanvas.Canvas // wraps gg.Context
fontSource *text.FontSource // font used to render per-frame stats
paletteView *wgpu.TextureView // handle to the color palette 1D texture view
fractalView gpucontext.TextureView // handle to the fractal texture view
fractalViewRelease func() // fractal TextureView release function
}
// renderer stores all objects required to render a fractal.
type renderer struct {
theme string // color scheme (green or red)
fractal *fractal // fractal being rendered
state *state
gpu *gpu
assets *assets
}
// uniforms stores per-frame uniforms.
type uniforms struct { // total: (2 uint32 + 14 float32) * 4 bytes == 64 bytes
paletteSize, powScale uint32
frameCount, maxIter float32 // block 1
width, height, scaleHi, scaleLo float32 // block 2
xRealHi, yImagHi, xRealLo, yImagLo float32 // block 3
cRealHi, cImagHi, cRealLo, cImagLo float32 // block 4
}
// ui encapsulates GUI and internal state and orchestrates animation and window lifecycle.
type ui struct {
app app // stores the GoGPU App instance
primaryWindow, // stores the primary window handle
aboutWindow window // stores the About window handle
renderer, // stores the renderer instance
animToken atomic.Value // stores animation state (active or paused)
lastFrameTime time.Time // timestamp recorded during the previous draw cycle; used to calculate FPS
initTokenOnce sync.Once // ensures the animation token is initialized (activated) exactly once
animating, // tracks persistent animation play / pause state across window visibility changes
pendingMenuRebuild, // flag used to signal that a new fractal was selected so the "Themes" menu should be rebuilt
hidePrimaryWindow, // defers hiding the primary window to OnUpdate() to avoid GoGPU mutex deadlock
hideAboutWindow, // defers hiding the About window to OnUpdate() to avoid GoGPU mutex deadlock
aboutWindowHasFocus atomic.Bool // tracks focus state of the About window so the correct window is closed when โŒ˜+W is pressed while both windows are visible
prog *progress.Progress // progress bar rendered to the terminal while a fractal is being rendered
progClose func() // stops the progress bar
}
fractal, theme, err := flags(flag.CommandLine, os.Args[1:])
if err != nil {
fmt.Fprintf(os.Stderr, "cannot parse flags: %v\n", err)
os.Exit(1)
}
// workaround for none of the following working as expected:
//
// gogpu.SetLogger(nil)
// gogpu.SetLogger(slog.New(slog.DiscardHandler))
// gogpu.SetLogger(slog.New(slog.NewTextHandler(io.Discard, nil)))
slog.SetDefault(slog.New(slog.DiscardHandler))
ctx, cancel := context.WithCancel(context.Background())
ui := new(ui)
ui.newProgressBar(ctx)
defer ui.progClose() // called via panic() paths (idempotent)
ui.app = &gogpuApp{App: gogpu.NewApp(gogpu.DefaultConfig().
WithAppName("Fractal").
WithTitle(fractal.titleText).
WithSize(mainWidth, mainHeight).
WithResizable(false))}
ui.app.SetQuitOnLastWindowClosed(false)
// trap SIGINT, SIGTERM, and SIGHUP to ensure clean progress
// bar shutdown so the hidden cursor is restored
sigChan := make(chan os.Signal, 1)
signal.Notify(sigChan,
syscall.SIGINT, // ctrl+c
syscall.SIGTERM, // standard kill signal
syscall.SIGHUP) // terminal closed, SSH disconnection, etc
go func() {
<-sigChan
ui.app.Quit() // triggers OnClose() to ensure clean progress bar shutdown
}()
shaderCode := map[string]string{
"mandelbrot": removeComments(commonShaderCode + "\n" + mandelbrotShaderCode),
"julia": removeComments(commonShaderCode + "\n" + juliaShaderCode),
}
ui.renderer.Store(newRenderer(fractal, shaderCode[fractal.kind], theme))
// GoGPU callback registrations and definitions
ui.app.OnSurfaceAvailable(func() {
ui.aboutWindow, err = ui.app.NewWindow(gogpu.DefaultConfig().
WithTitle("").
WithSize(aboutWidth, aboutHeight).
WithTransparent(true).
WithResizable(false).
WithHeaderAlignment(gogpu.HeaderAlignLeft)) // renders transparent title bar
if err != nil {
panic(err)
}
ui.aboutWindow.Hide()
ui.aboutWindowHasFocus.Store(false)
ui.primaryWindow = ui.app.PrimaryWindow()
ui.primaryWindow.SetOnKeyPress(func(key gpucontext.Key, mods gpucontext.Modifiers) {
if mods.HasSuper() && key == gpucontext.KeyW { // โŒ˜+W
if ui.aboutWindowHasFocus.Load() {
ui.hideAboutWin() // immediately hide the About window
} else if ui.primaryWindow.Visible() {
ui.hidePrimaryWin() // defer primaryWindow.Hide() to OnUpdate() to avoid GoGPU internal mutex deadlock
}
}
})
ui.primaryWindow.SetOnClose(func() bool {
if ui.primaryWindow.Visible() {
ui.hidePrimaryWin() // defer primaryWindow.Hide() to OnUpdate() to avoid GoGPU internal mutex deadlock
}
return false // reject the close window request and hide the window instead to preserve window handle and callbacks
})
// TODO(jbunds): fix the following bug:
//
// 1. launch the app
// 2. open the About window
// 3. use macOS "Hot Corners" to bring the primary application window into the foreground and give it focus
// 4. press โŒ˜+W
// 5. bug: the About window is closed instead of the primary application window
//
// i suspect this may be impossible to fix unless and until the GoGPU core API supports per-window focus tracking
ui.primaryWindow.SetOnPointer(func(e gpucontext.PointerEvent) {
if e.Type == gpucontext.PointerDown {
ui.aboutWindowHasFocus.Store(false)
}
})
cr := ui.renderer.Load().(*renderer)
cr.init(ui.app, cr.theme)
setAppMenu(ui)
addFractalsMenu(ctx, ui, shaderCode)
rebuildThemesMenu(ctx, ui, shaderCode) // seed the "Themes" menu's Action closures with the current renderer instance
// TODO(jbunds): fix whatever is removing the native "Window" menu,
// which may be triggered by customizing the menus per
// the call to app.SetMenu() or app.SetCustomMenu() (?)
//
// the "Window" menu only gets added without the calls to
// setAppMenu, addFractalsMenu, and rebuildThemesMenu
addWindowMenu(ui.app)
})
ui.app.EventSource().OnKeyPress(func(key gpucontext.Key, _ gpucontext.Modifiers) {
if key == gpucontext.KeySpace { // space bar
ui.toggleAnimation()
}
})
ui.app.OnUpdate(func(_ float64) {
if ui.pendingMenuRebuild.Swap(false) { // if a new fractal was selected
rebuildThemesMenu(ctx, ui, shaderCode) // rebuild the "Themes" menu to seed its Action closures with the new renderer instance
}
if ui.hidePrimaryWindow.Swap(false) { // if โŒ˜+W was pressed when the primary window had focus
ui.primaryWindow.Hide() // immediately hide the primary window
if ui.aboutWindow.Visible() { // if the About window is visible
ui.aboutWindow.Show() // focus the About window to enable the window control buttons
ui.aboutWindowHasFocus.Store(true) // record the About window now has focus
}
}
if ui.hideAboutWindow.Swap(false) { // if โŒ˜+W was pressed when the About window had focus
ui.hideAboutWin() // immediately hide the About window and focus the primary window if visible
}
})
ui.app.OnDraw(func(dc *gogpu.Context) {
ui.initTokenOnce.Do(func() {
ui.animToken.Store(tokenRef{token: ui.app.StartAnimation()})
ui.animating.Store(true)
})
cr := ui.renderer.Load().(*renderer)
cr.draw(dc, &ui.animToken)
elapsed := time.Since(ui.lastFrameTime).Milliseconds()
if elapsed > 0 {
cr.state.fps = float64(1000.0 / elapsed)
}
ui.lastFrameTime = time.Now()
if cr.state.frameCount <= maxPrecisionFrames {
ui.prog.Report(1, strconv.Itoa(cr.state.frameCount))
} else {
ui.prog.Close()
}
if ui.loadToken() != nil {
ui.app.RequestRedraw() // renders at VSync frequency (~60 FPS)
}
})
ui.app.OnClose(func() {
ui.renderer.Load().(*renderer).release()
gg.CloseAccelerator()
cancel()
ui.prog.Close()
})
ui.lastFrameTime = time.Now()
// main event loop
if err := ui.app.Run(); err != nil {
ui.renderer.Load().(*renderer).release()
gg.CloseAccelerator()
cancel()
panic(err) // deferred call to ui.prog.Close() triggered here
}
cancel() // unreachable since Run() blocks
}
// newProgressBar initializes a new progress bar, or replaces an active progress bar with a new one.
if u.progClose != nil { u.progClose() }
u.prog = progress.New(ctx, maxPrecisionFrames, os.Stderr,
progress.WithTracker(progress.Fraction),
progress.WithTheme("magma"),
)
u.progClose = u.prog.Close
}
// syncAnimation ensures the animation token matches the desired state:
// playing if `animating` is true and the primary window is visible; paused otherwise.
if u.animating.Load() && u.primaryWindow.Visible() {
if u.loadToken() == nil {
u.animToken.Store(tokenRef{token: u.app.StartAnimation()})
}
} else {
if token := u.loadToken(); token != nil {
token.Stop()
u.animToken.Store(tokenRef{})
}
}
}
// toggleAnimation toggles between pausing and resuming the animation loop,
// e.g., when the primary window is hidden, or the space bar is pressed.
if u.loadToken() == nil {
u.animating.Store(true)
} else {
u.animating.Store(false)
}
u.syncAnimation()
}
// hidePrimaryWin stops the active animation loop and defers hiding the primary window to OnUpdate().
if token := u.loadToken(); token != nil {
token.Stop()
}
u.hidePrimaryWindow.Store(true) // defer primaryWindow.Hide() to OnUpdate() to avoid GoGPU internal mutex deadlock
}
// hideAboutWin immediately hides the About window, updates the About window's
// focus-tracking flag, and focuses the primary window if visible.
u.aboutWindow.Hide() // immediately hide the About window
u.aboutWindowHasFocus.Store(false) // record the About window lost focus
if u.primaryWindow.Visible() { // if the primary window is not hidden
u.primaryWindow.Show() // focus the primary window to enable the window control buttons
}
}
// scheduleMenuRebuild marks the "Themes" menu for rebuild per the next OnUpdate() cycle.
u.pendingMenuRebuild.Store(true)
if u.primaryWindow.Visible() && u.loadToken() == nil {
u.app.RequestRedraw() // ensure OnUpdate() fires if animation is paused
}
}
// loadToken is an atomic.Value loading helper to reduce type assertion boilerplate.
r, _ := u.animToken.Load().(tokenRef)
return r.token
}
// newRenderer constructs and returns the *renderer used to store renderer state.
xRealHi, xRealLo := splitFloat64(fractal.params.xReal)
yImagHi, yImagLo := splitFloat64(fractal.params.yImag)
cRealHi, cRealLo := splitFloat64(fractal.params.cReal)
cImagHi, cImagLo := splitFloat64(fractal.params.cImag)
return &renderer{
fractal: fractal,
theme: theme,
gpu: &gpu{shaderCode: shaderCode},
assets: &assets{},
state: &state{
viewportWidth: viewportWidth,
xRealHi: xRealHi,
xRealLo: xRealLo,
yImagHi: yImagHi,
yImagLo: yImagLo,
cRealHi: cRealHi,
cRealLo: cRealLo,
cImagHi: cImagHi,
cImagLo: cImagLo,
},
}
}
// init initializes all resources required to render frames in the main application window.
// TODO(jbunds): consider using a sync.Map-based main-global (singleton) cache to cache all static resources
// (arguably overkill as the overhead of reinstantiating static resources is not a bottleneck)
//
// func getOrInit[T any](r *renderer, key string, load func() (T, error)) T {
// if val, ok := r.cache.Load(key); ok {
// return val.(T)
// }
// onceKey := key + ":_once"
// val, _ := r.cache.LoadOrStore(onceKey, new(sync.Once))
// once := val.(*sync.Once)
// var result T
// var onceErr error
// once.Do(func() {
// result, onceErr = load()
// if onceErr != nil { panic(onceErr) }
// r.cache.Store(key, result)
// r.cache.Delete(onceKey)
// })
// val, ok := r.cache.Load(key)
// if !ok { panic("cache invariant broken: value not stored") }
// return val.(T)
// }
var err error
r.assets.fontSource, err = loadFont(os.DirFS("/").(fs.StatFS))
if err != nil {
panic(err)
}
r.gpu.device = app.DeviceProvider().Device()
r.state.paletteColors,
r.gpu.paletteTex,
r.gpu.uniformBuf,
r.gpu.bgLayout0,
r.gpu.bgLayout1,
r.gpu.pipeline = initResources(r.gpu.device, r.gpu.shaderCode, theme)
r.assets.canvas, err = ggcanvas.New(app.GPUContextProvider(), mainWidth, mainHeight)
if err != nil {
panic(err)
}
r.assets.canvas.Context().SetFont(r.assets.fontSource.Face(12))
r.assets.paletteView, err = r.gpu.device.CreateTextureView(r.gpu.paletteTex, &wgpu.TextureViewDescriptor{
Format: gputypes.TextureFormatR32Uint,
Dimension: gputypes.TextureViewDimension1D,
})
if err != nil {
panic(err)
}
err = r.gpu.device.Queue().WriteTexture(
&wgpu.ImageCopyTexture{Texture: r.gpu.paletteTex},
unsafe.Slice(
(*byte)(unsafe.Pointer(&r.state.paletteColors[0])), // #nosec G103 - audited
len(r.state.paletteColors) * 4),
&wgpu.ImageDataLayout{BytesPerRow: paletteSize * 4},
&wgpu.Extent3D{Width: paletteSize, Height: 1, DepthOrArrayLayers: 1})
if err != nil {
panic(err)
}
r.gpu.staticBindGroup, err = r.gpu.device.CreateBindGroup(&wgpu.BindGroupDescriptor{
Layout: r.gpu.bgLayout0,
Entries: []wgpu.BindGroupEntry{
{Binding: 0, Size: 64, Buffer: r.gpu.uniformBuf},
{Binding: 1, TextureView: r.assets.paletteView},
},
})
if err != nil {
panic(err)
}
r.assets.fractalView,
r.assets.fractalViewRelease = r.assets.canvas.Context().CreateOffscreenTexture(mainWidth, mainHeight)
if r.assets.fractalViewRelease == nil {
fmt.Fprint(os.Stderr, "GPU unavailable") // TODO(jbunds): implement CPU fallback ?
}
}
// draw renders a new frame to the canvas.
if r.assets.canvas.Context() == nil {
return
}
if ref, _ := token.Load().(tokenRef); ref.token != nil {
if r.state.frameCount > maxPrecisionFrames { // precision exhausted (perhaps many frames ago, depending on the specific fractal)
ref.token.Stop()
return
}
r.state.viewportWidth *= scaleFactor
r.state.frameCount++
}
// per-frame state updates
unis := updateUniforms( // magnification logic
r.state.frameCount, r.fractal.params.powScale,
r.state.xRealHi, r.state.xRealLo,
r.state.yImagHi, r.state.yImagLo,
r.state.cRealHi, r.state.cRealLo,
r.state.cImagHi, r.state.cImagLo,
r.state.viewportWidth, r.fractal.params.maxIter(r.state.frameCount),
)
err := r.gpu.device.Queue().WriteBuffer(
r.gpu.uniformBuf, 0,
unsafe.Slice((*byte)(unsafe.Pointer(unis)), 64)) // #nosec G103 - audited
if err != nil {
panic(err)
}
fractalViewBindGroup := r.fractalViewBindGroup()
defer fractalViewBindGroup.Release()
encoder, err := r.gpu.device.CreateCommandEncoder(nil)
if err != nil {
panic(err)
}
pass, err := encoder.BeginComputePass(nil)
if err != nil {
panic(err)
}
pass.SetPipeline(r.gpu.pipeline)
pass.SetBindGroup(0, r.gpu.staticBindGroup, nil)
pass.SetBindGroup(1, fractalViewBindGroup, nil)
width, height := dc.SurfaceSize() // https://pkg.go.dev/github.com/gogpu/gogpu#App.ScaleFactor
pass.Dispatch(((width + 15) / 16), ((height + 7) / 8), 1)
err = pass.End()
if err != nil {
panic(err)
}
cmds, err := encoder.Finish()
if err != nil {
panic(err)
}
_, err = r.gpu.device.Queue().Submit(cmds)
if err != nil {
panic(err)
}
r.drawStats()
err = r.assets.canvas.RenderDirect(dc.RenderTarget().SurfaceView(), width, height)
if err != nil {
fmt.Fprintln(os.Stderr, err)
}
}
// drawStatus draws a rectangular box in the bottom-left corner
// of the main window showing some basic runtime stats.
err := r.assets.canvas.Draw(func(cc *gg.Context) {
cc.DrawGPUTextureBase(r.assets.fractalView, 0, 0, mainWidth, mainHeight)
cc.SetRGBA(0, 0, 0, 0.15)
cc.DrawRoundedRectangle(10, mainHeight - 40, 336, 30, 4)
cc.Fill()
cc.SetColor(gg.Cyan)
cc.DrawString(fmt.Sprintf("FPS: %.0f", r.state.fps ), 18, mainHeight - 20)
cc.SetColor(gg.Green)
cc.DrawString(fmt.Sprintf("magnification: %e", 1 / r.state.viewportWidth), 72, mainHeight - 20)
cc.SetColor(gg.Cyan)
cc.DrawString(fmt.Sprintf("frames: %d", r.state.frameCount ), 258, mainHeight - 20)
})
if err != nil {
panic(err)
}
}
// fractalViewBindGroup creates the BindGroup holding the per-frame TextureView of the rendered fractal.
fractalViewBindGroup, err := r.gpu.device.CreateBindGroup(&wgpu.BindGroupDescriptor{
Layout: r.gpu.bgLayout1,
Entries: []wgpu.BindGroupEntry{{
Binding: 0,
TextureView: (*wgpu.TextureView)(r.assets.fractalView.Pointer()),
}},
})
if err != nil {
panic(err)
}
return fractalViewBindGroup
}
// release marks resources for deallocation.
r.assets.fractalViewRelease()
r.assets.canvas.Close()
r.assets.fontSource.Close()
r.assets.paletteView.Release()
r.gpu.pipeline.Release()
r.gpu.bgLayout0.Release()
r.gpu.bgLayout1.Release()
r.gpu.paletteTex.Release()
r.gpu.uniformBuf.Release()
r.gpu.staticBindGroup.Release()
}
// updateUniforms updates the per-frame uniforms passed to the GPU shader.
frameCount int,
powScale uint32,
xRealHi, xRealLo,
yImagHi, yImagLo,
cRealHi, cRealLo,
cImagHi, cImagLo float32,
viewportWidth, maxIter float64) *uniforms {
scaleHi, scaleLo := splitFloat64(viewportWidth)
return &uniforms{
powScale: powScale,
paletteSize: paletteSize,
scaleHi: scaleHi,
scaleLo: scaleLo,
width: float32(mainWidth),
height: float32(mainHeight),
maxIter: float32(maxIter),
frameCount: float32(frameCount),
xRealHi: xRealHi,
xRealLo: xRealLo,
yImagHi: yImagHi,
yImagLo: yImagLo,
cRealHi: cRealHi,
cRealLo: cRealLo,
cImagHi: cImagHi,
cImagLo: cImagLo,
}
}
// removeComments removes comment-only lines from WGSL shader code.
commentLine := regexp.MustCompile(`(?m)^\s*//.*\n`)
return strings.TrimSpace(commentLine.ReplaceAllString(text, ""))
}
// splitFloat64 splits a float64 into two float32s.
high := float32(v)
low := float32(v - float64(high))
return high, low
}