ParticleSketch.kt

@file:Suppress("MemberVisibilityCanBePrivate")

package afield.simulation.particles

import afield.document.context.fromContext
import afield.document.graph.Box
import afield.graphics.color.rgba
import afield.graphics.render.ComputeMesh
import afield.graphics.render.ComputeModulePointer
import afield.graphics.render.Mesh
import afield.graphics.render.Module
import afield.graphics.render.ModulePointer
import afield.graphics.render.computeMemoryBarrier
import afield.graphics.render.fullscreenTriangle
import afield.graphics.render.graphicsMemoryBarrier
import afield.graphics.render.parameters
import afield.graphics.render.pointer
import afield.graphics.shader.fromViews
import afield.graphics.texture.TextureSampler
import afield.graphics.vulkan.Renderer
import afield.graphics.vulkan.Target
import afield.graphics.vulkan.buffers.FrameBufferGroup
import afield.graphics.vulkan.buffers.defaultBuffer
import afield.graphics.vulkan.device.CommandBuffer
import afield.graphics.vulkan.targets.PointAccumulation
import afield.graphics.vulkan.targets.PointBlend
import afield.graphics.vulkan.targets.PointShaders
import afield.graphics.vulkan.util.GpuNeighbourGrid
import afield.graphics.vulkan.util.WindowPointer
import afield.math.linalg.Rectanglef
import afield.math.linalg.Vec4
import afield.math.linalg.Vector2f
import afield.math.linalg.vec
import afield.ui.output.controlGroup
import afield.ui.output.displayOpen
import org.lwjgl.system.MemoryStack
import org.lwjgl.vulkan.VK13.VK_ACCESS_2_SHADER_READ_BIT
import org.lwjgl.vulkan.VK13.VK_ACCESS_2_SHADER_WRITE_BIT
import org.lwjgl.vulkan.VK13.VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT
import org.lwjgl.vulkan.VK13.VK_PIPELINE_STAGE_2_FRAGMENT_SHADER_BIT
import org.lwjgl.vulkan.VK13.VK_PIPELINE_STAGE_2_VERTEX_SHADER_BIT
import org.lwjgl.vulkan.VK13.VK_STRUCTURE_TYPE_DEPENDENCY_INFO
import org.lwjgl.vulkan.VK13.VK_STRUCTURE_TYPE_MEMORY_BARRIER_2
import org.lwjgl.vulkan.VK13.vkCmdPipelineBarrier2
import org.lwjgl.vulkan.VkDependencyInfo
import org.lwjgl.vulkan.VkMemoryBarrier2
import java.nio.ByteOrder
import kotlin.math.ceil
import kotlin.math.floor
import kotlin.math.max
import kotlin.math.min
import kotlin.math.pow
import kotlin.math.sqrt

    private class Resources(val renderer: Renderer, val shape: Shape) : Target {
        var settings = ParticleSketchSettings()
        var owner: ComputeModulePointer? = null

        val particles: List<FrameBufferGroup> = List(2) { i ->
            defaultBuffer(name = "${shape.name}-particles-$i") {
                size = shape.capacity.toLong() * PARTICLE_FLOATS * 4
            }.toBuffer(renderer)
        }

        val strokes: FrameBufferGroup = defaultBuffer(name = "${shape.name}-strokes") {
            size = 16L + shape.maxStrokes.toLong() * STROKE_BYTES
        }.toBuffer(renderer)

        // Written by the CPU every frame, so one per frame in flight.
        val world: FrameBufferGroup = defaultBuffer(name = "${shape.name}-world") {
            size = WORLD_BYTES.toLong()
            perFrame = true
        }.toBuffer(renderer)

        val spawns: FrameBufferGroup = defaultBuffer(name = "${shape.name}-spawns") {
            size = SPAWNS_PER_FRAME * 8L
            perFrame = true
        }.toBuffer(renderer)

        val grid = GpuNeighbourGrid(
            capacity = shape.capacity,
            cellsPerSide = CELLS,
            domainMin = Vector2f(-GRID_MARGIN, -GRID_MARGIN),
            domainSize = CELLS * CELL_SIZE,
            renderer = renderer,
            name = "${shape.name}-grid",
        )

        val accumulation = PointAccumulation(renderer) {
            name = "${shape.name}-picture"
            blend = PointBlend.ALPHA
            fade = 1.0
        }

        private val gridTarget = grid.target({ particles[current] }, { count() }, PARTICLE_FLOATS, 0)

        @Volatile var current = 0
        @Volatile var spawned = 0L
        @Volatile var time = 0.0
        var phase = 0
        private var frame = 0
        private var lastNanos = 0L
        @Volatile private var resetPending = false

        val queue = ArrayDeque<Float>()

        // the pointer and the window's size in points, read on the window's poll
        private val brush = SpawnBrush()
        val windowPointer = WindowPointer(renderer.window, "${shape.name}-pointer") { p ->
            brush.spacing = settings.paintSpacing
            brush.step(p.x.toDouble(), p.y.toDouble(), p.down && settings.paintWithMouse, p.inside, ::enqueue)
        }
        val worldWidth get() = windowPointer.width
        val worldHeight get() = windowPointer.height

        private var module: ComputeModulePointer? = null
        private var mesh: ComputeMesh? = null
        private val bounds = Rectanglef(0f, 0f, 0f, 0f)
        private val sampler = TextureSampler(renderer, 1, false)

        init {
            buildDrawing()
        }

        /** Particles alive; 0 as soon as [reset] is asked for, though the frame applies it. */
        fun count(): Int = if (resetPending) 0 else min(spawned, shape.capacity.toLong()).toInt()

        fun install(module: ComputeModulePointer, mesh: ComputeMesh) {
            this.module = module
            this.mesh = mesh
        }

        fun reset() {
            synchronized(queue) { queue.clear() }
            resetPending = true
            accumulation.clear()
        }

        fun enqueue(x: Double, y: Double) {
            synchronized(queue) {
                queue.addLast(x.toFloat())
                queue.addLast(y.toFloat())
            }
        }

        fun paintAlong(x0: Double, y0: Double, x1: Double, y1: Double) =
            SpawnBrush.along(x0, y0, x1, y1, settings.paintSpacing, ::enqueue)

        private lateinit var strokesUnder: Mesh
        private lateinit var dots: Mesh
        private lateinit var strokesOver: Mesh

        private fun buildDrawing() {
            strokesUnder = strokeLayer("${shape.name}-links", 0f)
            dots = dotLayer("${shape.name}-dots")
            strokesOver = strokeLayer("${shape.name}-overlays", 1f)
        }

        private fun targetSize() = vec(accumulation.width.toFloat(), accumulation.height.toFloat())

        private fun strokeLayer(name: String, layer: Float): Mesh {
            val m = Module(accumulation.pointParameters(name) {
                vertex = ParticleShaders.source("strokes_vertex.glsl")
                fragment = ParticleShaders.source("shape_fragment.glsl")
            }).pointer()
            val mesh = Mesh(renderer, m).resize(shape.maxStrokes * 6)
            mesh.push.bind("strokes") { strokes }
            mesh.push.bind("size") { vec(worldWidth, worldHeight) }
            mesh.push.bind("targetSize") { targetSize() }
            mesh.push.bind("capacity") { shape.maxStrokes }
            mesh.push.bind("layer") { layer }
            m += "$name-quads" to mesh
            accumulation.add(name, m)
            return mesh
        }

        private fun dotLayer(name: String): Mesh {
            val m = Module(accumulation.pointParameters(name) {
                vertex = ParticleShaders.source("dots_vertex.glsl")
                fragment = ParticleShaders.source("shape_fragment.glsl")
            }).pointer()
            val mesh = Mesh(renderer, m).resize(shape.capacity * 6)
            mesh.push.bind("particles") { particles[current] }
            mesh.push.bind("size") { vec(worldWidth, worldHeight) }
            mesh.push.bind("targetSize") { targetSize() }
            mesh.push.bind("count") { count() }
            m += "$name-quads" to mesh
            accumulation.add(name, m)
            return mesh
        }

        private val displayModule: ModulePointer by lazy { buildDisplay() }

        /**
         * Make this the sketch running under its name: its frame and its display. Resources kept
         * for a re-run are activated again, because another shape of sketch under the same name
         * may have taken both in between.
         */
        fun activate() {
            renderer.renderList.preTargets["${shape.name}-sketch"] = this
            renderer.renderList += "${shape.name}-display" to displayModule
        }

        private fun buildDisplay(): ModulePointer {
            val name = "${shape.name}-display"
            val m = Module(
                parameters(renderer) {
                    this.name = name
                    depthTest = false
                    depthWrite = false
                    blend = false
                    vertex = PointShaders.source("fullscreen_vertex.glsl")
                    fragment = ParticleShaders.source("display_fragment.glsl")
                    textureResolver = { tex ->
                        if (tex.name == "accumulation") fromViews(listOf(accumulation.latestImageView), tex.name, sampler) else null
                    }
                }
            ).pointer()
            m += "$name-quad" to fullscreenTriangle(m).apply {
                push.bind("background") { settings.background }
            }
            return m
        }

        /** Takes the spawns queued since the last frame, at most a frame's worth. */
        private fun drainQueue(): FloatArray = synchronized(queue) {
            val n = min(queue.size / 2, SPAWNS_PER_FRAME)
            FloatArray(n * 2) { queue.removeFirst() }
        }

        override fun render(command: CommandBuffer, index: Int) {
            val module = module ?: return
            val mesh = mesh ?: return
            // the world is the window's size in points, which the first poll brings
            if (!windowPointer.polled) return
            val s = settings

            if (resetPending) {
                spawned = 0L
                time = 0.0
                resetPending = false
            }

            val now = System.nanoTime()
            // Between steps nothing is recorded at all: the picture, and so the window, hold still.
            if (s.fps > 0 && lastNanos != 0L && now - lastNanos < (1e9 / s.fps).toLong()) return
            val frameSeconds = if (lastNanos == 0L) 1.0 / 60.0 else ((now - lastNanos) / 1e9).coerceIn(0.0, 1.0)
            val dt = frameSeconds.coerceAtMost(MAX_DT)
            lastNanos = now
            val stepping = !s.paused
            if (stepping) time += dt
            frame++

            var batch = drainQueue()
            // More paint in one frame than there is room for: only the newest capacity's worth can
            // survive, and two spawns writing one slot in the same dispatch would race, so the
            // older ones are skipped (they still count, so ids and the ring stay as if painted).
            val overflow = batch.size / 2 - shape.capacity
            if (overflow > 0) {
                spawned += overflow
                batch = batch.copyOfRange(overflow * 2, batch.size)
            }
            val spawnCount = batch.size / 2
            val spawnBase = (spawned % shape.capacity).toInt()
            val spawnId = spawned + 1
            spawned += spawnCount
            val alive = count()

            spawns.mapAndWrite { put(batch) }
            world.mapAndWrite_bytes {
                order(ByteOrder.nativeOrder())
                val c = s.particleColor
                putFloat(c.x); putFloat(c.y); putFloat(c.z); putFloat(c.w)
                putFloat(worldWidth); putFloat(worldHeight)
                putFloat(windowPointer.x); putFloat(windowPointer.y)
                putFloat(s.gravity.x); putFloat(s.gravity.y)
                putFloat(-GRID_MARGIN); putFloat(-GRID_MARGIN)
                putFloat(time.toFloat())
                putFloat(if (stepping) dt.toFloat() else 0f)
                putFloat(s.damping.toFloat())
                putFloat(s.maxSpeed.toFloat())
                putFloat(s.pairRadius.toFloat())
                putFloat(s.particleRadius.toFloat())
                putFloat(1f / CELL_SIZE)
                putInt(CELLS)
                putInt(if (windowPointer.inside) 1 else 0)
                putInt(if (windowPointer.down) 1 else 0)
                putInt(alive)
                putInt(shape.capacity)
                putInt(spawnBase)
                putInt(spawnCount)
                putInt(spawnId.toInt())
                putInt(shape.maxStrokes)
                putInt(frame)
                putInt(s.seed)
            }

            // last frame's draws read the buffers this frame's passes write
            graphicsToComputeBarrier(command)

            // 1. spawn: the paint queue becomes particles; the stroke list is cleared; draw() runs
            phase = 0
            mesh.x = ceil(max(spawnCount, 1) / LOCAL_SIZE.toDouble()).toInt()
            module.render(command, bounds)
            computeMemoryBarrier(command)

            if (stepping && alive > 0) {
                // 2. hash: the neighbour grid over every particle (it ends with its own barrier)
                gridTarget.render(command, index)
                // 3. step: eachPair, update, integrate -- from one buffer into the other
                phase = 1
                mesh.x = ceil(alive / LOCAL_SIZE.toDouble()).toInt()
                module.render(command, bounds)
                current = 1 - current
            }
            graphicsMemoryBarrier(command)

            // 4. draw: fade, links, dots, overlays
            // madp fades by 1 - trail once per frame at 60 frames a second; scaled to this step's
            // length so the trail is as long in seconds whatever fps is. Wall-clock time, so a
            // paused picture fades as madp's does.
            val keep = s.trail.coerceIn(0.0, 1.0).pow(frameSeconds * 60.0)
            accumulation.settings.fade = (1.0 - keep).coerceIn(0.0, 1.0)
            accumulation.render(command, index)
        }

        private fun graphicsToComputeBarrier(command: CommandBuffer) {
            MemoryStack.stackPush().use { stack ->
                val m = VkMemoryBarrier2.calloc(1, stack)
                m[0]
                    .sType(VK_STRUCTURE_TYPE_MEMORY_BARRIER_2)
                    .srcStageMask(VK_PIPELINE_STAGE_2_VERTEX_SHADER_BIT or VK_PIPELINE_STAGE_2_FRAGMENT_SHADER_BIT)
                    .srcAccessMask(VK_ACCESS_2_SHADER_READ_BIT)
                    .dstStageMask(VK_PIPELINE_STAGE_2_COMPUTE_SHADER_BIT)
                    .dstAccessMask(VK_ACCESS_2_SHADER_READ_BIT or VK_ACCESS_2_SHADER_WRITE_BIT)
                val dep = VkDependencyInfo.calloc(stack).sType(VK_STRUCTURE_TYPE_DEPENDENCY_INFO).pMemoryBarriers(m)
                vkCmdPipelineBarrier2(command.vkCommandBuffer, dep)
            }
        }
    }
}