alife_fire.kt

data class ShaderForestFireSettings(
    /** Prefix for module, mesh and pre-target names. Two systems in one renderer need two names. */
    var name: String = "forest-fire",
    /** Forest width in texels. */
    var width: Int = 1024,
    /** Forest height in texels. */
    var height: Int = 1024,
    /** Format of all three textures. */
    var format: Int = VK_FORMAT_R16G16B16A16_SFLOAT,
    /**
     * Clear colour of every texture. Transparent black, because alpha is not opacity here: it is the
     * fourth kind of fire, and an opaque clear would light all of it at once.
     */
    var clearColor: Vec4 = vec(0, 0, 0, 0),
    /** The fraction of texels that are trees. Below about 0.4 fire stops crossing the gaps. */
    var treeDensity: Float = 0.72f,
    /** Size in texels of the coarse cells the forest is hashed on, so trees grow in clumps. */
    var treeScalePx: Float = 3f,
    /** How much per-texel noise is mixed into the clumps, 0 to 1. */
    var treeJitter: Float = 0.28f,
    /** How far, in texels, a tree looks for burning neighbours. */
    var spreadStepPx: Float = 1f,
    /** Gain on a burning neighbour's fire before it is compared with [ignitionThreshold]. */
    var spreadGain: Float = 1.02f,
    /**
     * The chance, each frame, that one burning neighbour lights a ready tree. At 1 every front
     * moves one texel a frame in all eight directions and comes out square; below 1 it moves by
     * chance, slower on average, and rounds itself into a ragged circle.
     */
    var spreadProbability: Float = 0.5f,
    /** Gain on the ignition texture before it is compared with [ignitionThreshold]. */
    var paintGain: Float = 1.0f,
    /** The signal at or above which a ready tree catches. */
    var ignitionThreshold: Float = 0.33f,
    /** The fire level at or above which a tree is burning. */
    var activeFireThreshold: Float = 0.82f,
    /** The fire level at or below which a tree can catch again. */
    var readyThreshold: Float = 0.03f,
    /**
     * How fast each kind's fire level falls per frame. With the default thresholds red burns for 18
     * frames and is ready again after 97; alpha takes two and a half times as long.
     */
    var refractoryDecayPerFrame: FireChannels = FireChannels(0.010f, 0.008f, 0.006f, 0.004f),
    /** How fast the ignition texture fades per frame, per kind. */
    var ignitionDecayPerFrame: FireChannels = FireChannels(0.002f, 0.002f, 0.002f, 0.002f),
    /** How fast the ink fades per frame, per channel. */
    var inkDecayPerFrame: FireChannels = FireChannels(0.004f, 0.004f, 0.004f, 0.004f),
    /** The fraction of a burning neighbour's ink a burning tree takes on. */
    var inkSpreadGain: Float = 0.92f,
    /** Seed for the forest. */
    var randomSeed: Long = 1L
)

    init {
        require(settings.width > 0) { "width must be positive" }
        require(settings.height > 0) { "height must be positive" }
        require(settings.treeScalePx > 0f) { "treeScalePx must be positive" }

        fireInputPreserveModule = fullscreenPass(fireInput, "fire-input-preserve", PRESERVE_FRAGMENT, blend = false) { tex ->
            if (tex.name == "sourceTex") fromViews(listOf(fireInput.otherImageView), tex.name, fireInputSampler) else null
        }
        fireInputPreserveModule.quad {
            push.bind("decay") { settings.ignitionDecayPerFrame.toVec4() }
            push.bind("clearRequested") { if (clearFireInputRequested) 1f else 0f }
        }
        fireInput += "preserve" to fireInputPreserveModule

        // the state reads its own previous frame and the ignition this frame has just drawn
        stateUpdateModule = fullscreenPass(state, "state-update", STATE_UPDATE_FRAGMENT, blend = false) { tex ->
            when (tex.name) {
                "sourceTex" -> fromViews(listOf(state.otherImageView), tex.name, stateSampler)
                "ignitionTex" -> fromViews(listOf(fireInput.latestImageView), tex.name, fireInputSampler)
                else -> null
            }
        }
        stateUpdateModule.quad {
            push.bind("refractoryDecay") { settings.refractoryDecayPerFrame.toVec4() }
            push.bind("spreadStepPx") { settings.spreadStepPx }
            push.bind("spreadGain") { settings.spreadGain }
            push.bind("paintGain") { settings.paintGain }
            push.bind("ignitionThreshold") { settings.ignitionThreshold }
            push.bind("activeFireThreshold") { settings.activeFireThreshold }
            push.bind("readyThreshold") { settings.readyThreshold }
            push.bind("treeDensity") { settings.treeDensity }
            push.bind("treeScalePx") { settings.treeScalePx }
            push.bind("treeJitter") { settings.treeJitter }
            push.bind("seed") { settings.randomSeed.toInt() }
            push.bind("clearRequested") { if (clearStateRequested) 1f else 0f }
            push.bind("spreadProbability") { settings.spreadProbability }
            push.bind("frame") { frame.toInt() }
        }
        state += "update" to stateUpdateModule

        inkPreserveModule = fullscreenPass(ink, "ink-preserve", PRESERVE_FRAGMENT, blend = false) { tex ->
            if (tex.name == "sourceTex") fromViews(listOf(ink.otherImageView), tex.name, inkSampler) else null
        }
        inkPreserveModule.quad {
            push.bind("decay") { settings.inkDecayPerFrame.toVec4() }
            push.bind("clearRequested") { if (clearInkRequested) 1f else 0f }
        }
        ink += "preserve" to inkPreserveModule

        inkPropagateModule = fullscreenPass(ink, "ink-propagate", INK_PROPAGATE_FRAGMENT, blend = true) { tex ->
            when (tex.name) {
                "sourceTex" -> fromViews(listOf(ink.otherImageView), tex.name, inkSampler)
                "stateTex" -> fromViews(listOf(state.latestImageView), tex.name, stateSampler)
                else -> null
            }
        }
        inkPropagateModule.quad {
            push.bind("decay") { settings.inkDecayPerFrame.toVec4() }
            push.bind("activeFireThreshold") { settings.activeFireThreshold }
            push.bind("spreadStepPx") { settings.spreadStepPx }
            push.bind("inkSpreadGain") { settings.inkSpreadGain }
        }
        ink += "propagate" to inkPropagateModule

        Box.context.get()?.attachShader(stateUpdateModule)
    }

    private fun pingPong() = PingPongRenderToTexture(renderer) {
        width = settings.width
        height = settings.height
        format = settings.format
        colorFormats = listOf(settings.format)
        clearColor = settings.clearColor
        colorStoreOp = VK_ATTACHMENT_STORE_OP_STORE
        colorSampled = true
        useDepth = false
    }

    private fun Parameters.additive() {
        blend = true
        srcColor = VK_BLEND_FACTOR_ONE
        dstColor = VK_BLEND_FACTOR_ONE
        colorBlendOp = VK_BLEND_OP_ADD
        srcAlpha = VK_BLEND_FACTOR_ONE
        dstAlpha = VK_BLEND_FACTOR_ONE
        alphaBlendOp = VK_BLEND_OP_ADD
    }

    private fun fullscreenPass(
        target: PingPongRenderToTexture,
        suffix: String,
        fragmentSource: String,
        blend: Boolean,
        resolver: (Introspected.Texture) -> Pair<HasImageViewArray, TextureSampler>?
    ): ModulePointer = Module(
        parameters(renderer) {
            +target
            name = "${settings.name}-$suffix"
            depthTest = false
            if (blend) additive() else this.blend = false
            vertex = FULLSCREEN_TRIANGLE_VERTEX
            fragment = fragmentSource
            textureResolver = resolver
        }
    ).pointer()

    private fun additiveParameters(
        target: PingPongRenderToTexture,
        name: String,
        configure: Parameters.() -> Unit,
        resolver: (Introspected.Texture) -> Pair<HasImageViewArray, TextureSampler>?
    ): Parameters = parameters(renderer) {
        +target
        this.name = name
        depthTest = false
        additive()
        textureResolver = resolver
        configure()
    }

    private fun ModulePointer.quad(name: String = current.parameters.name, bind: Mesh.() -> Unit) {
        current.meshes["$name-quad"] = fullscreenTriangle(this).apply(bind)
    }

    companion object {
        private val FULLSCREEN_TRIANGLE_VERTEX = """
            #version 450

            vec2 fullscreenVertex(uint index) {
                if (index == 0u) return vec2(-1.0, -1.0);
                if (index == 1u) return vec2(3.0, -1.0);
                return vec2(-1.0, 3.0);
            }

            void main() {
                gl_Position = vec4(fullscreenVertex(uint(gl_VertexIndex)), 0.0, 1.0);
            }
        """.trimIndent()

        private val PRESERVE_FRAGMENT = """
            #version 450

            layout(set = 0, binding = 0) uniform sampler2D sourceTex;

            layout(push_constant) uniform Push {
                vec4 decay;
                float clearRequested;
            } push;

            layout(location = 0) out vec4 outColor;

            void main() {
                ivec2 dims = textureSize(sourceTex, 0);
                vec4 previous = texelFetch(sourceTex, clamp(ivec2(gl_FragCoord.xy), ivec2(0), dims - 1), 0);
                outColor = push.clearRequested > 0.5 ? vec4(0.0) : clamp(previous - push.decay, vec4(0.0), vec4(1.0));
            }
        """.trimIndent()

        private val INK_PROPAGATE_FRAGMENT = """
            #version 450

            layout(set = 0, binding = 0) uniform sampler2D sourceTex;
            layout(set = 0, binding = 1) uniform sampler2D stateTex;

            layout(push_constant) uniform Push {
                vec4 decay;
                float activeFireThreshold;
                float spreadStepPx;
                float inkSpreadGain;
            } push;

            layout(location = 0) out vec4 outColor;

            float max4(vec4 v) {
                return max(max(v.x, v.y), max(v.z, v.w));
            }

            float burning(ivec2 p) {
                float denom = max(1.0 - push.activeFireThreshold, 0.0001);
                return clamp((max4(texelFetch(stateTex, p, 0)) - push.activeFireThreshold) / denom, 0.0, 1.0);
            }

            vec4 inkFrom(ivec2 pixel, ivec2 dims, float weight) {
                ivec2 p = clamp(pixel, ivec2(0), dims - 1);
                return texelFetch(sourceTex, p, 0) * (burning(p) * weight);
            }

            // added to the preserve pass's output, so the result is the larger of the faded ink here
            // and the strongest ink handed on by a burning neighbour
            void main() {
                ivec2 dims = textureSize(sourceTex, 0);
                ivec2 pixel = clamp(ivec2(gl_FragCoord.xy), ivec2(0), dims - 1);
                if (burning(pixel) <= 0.0) {
                    outColor = vec4(0.0);
                    return;
                }

                vec4 previous = texelFetch(sourceTex, pixel, 0);
                vec4 decayed = max(previous - push.decay, vec4(0.0));

                int s = max(int(round(push.spreadStepPx)), 1);
                float axial = push.inkSpreadGain;
                float diagonal = push.inkSpreadGain * 0.7071;
                vec4 carried = previous;
                carried = max(carried, inkFrom(pixel + ivec2(s, 0), dims, axial));
                carried = max(carried, inkFrom(pixel + ivec2(-s, 0), dims, axial));
                carried = max(carried, inkFrom(pixel + ivec2(0, s), dims, axial));
                carried = max(carried, inkFrom(pixel + ivec2(0, -s), dims, axial));
                carried = max(carried, inkFrom(pixel + ivec2(s, s), dims, diagonal));
                carried = max(carried, inkFrom(pixel + ivec2(-s, s), dims, diagonal));
                carried = max(carried, inkFrom(pixel + ivec2(s, -s), dims, diagonal));
                carried = max(carried, inkFrom(pixel + ivec2(-s, -s), dims, diagonal));

                outColor = max(carried - decayed, vec4(0.0));
            }
        """.trimIndent()

        private val PREVIEW_VERTEX = """
            #version 450

            layout(location = 0) out vec2 uv;

            void main() {
                vec2 p = gl_VertexIndex == 0 ? vec2(-1.0, -1.0) : (gl_VertexIndex == 1 ? vec2(3.0, -1.0) : vec2(-1.0, 3.0));
                uv = p * 0.5 + 0.5;
                gl_Position = vec4(p, 0.0, 1.0);
            }
        """.trimIndent()

        private val PREVIEW_FRAGMENT = """
            #version 450

            layout(set = 0, binding = 0) uniform sampler2D stateTex;
            layout(set = 0, binding = 1) uniform sampler2D inkTex;

            layout(push_constant) uniform Push {
                float activeFireThreshold;
                float readyThreshold;
            } push;

            layout(location = 0) in vec2 uv;
            layout(location = 0) out vec4 outColor;

            float max4(vec4 v) {
                return max(max(v.x, v.y), max(v.z, v.w));
            }

            vec3 fireKindColor(vec4 s) {
                return s.r * vec3(1.0, 0.18, 0.05) + s.g * vec3(0.15, 1.0, 0.18)
                    + s.b * vec3(0.15, 0.45, 1.0) + s.a * vec3(1.0, 0.2, 0.95);
            }

            void main() {
                ivec2 dims = textureSize(stateTex, 0);
                ivec2 pixel = clamp(ivec2(uv * vec2(dims)), ivec2(0), dims - 1);
                vec4 state = texelFetch(stateTex, pixel, 0);
                vec4 ink = texelFetch(inkTex, pixel, 0);

                float level = max4(state);
                float burning = clamp((level - push.activeFireThreshold) / max(1.0 - push.activeFireThreshold, 0.0001), 0.0, 1.0);
                float recovering = clamp((level - push.readyThreshold) / max(push.activeFireThreshold - push.readyThreshold, 0.0001), 0.0, 1.0) * (1.0 - burning);

                vec3 base = fireKindColor(state);
                vec3 inkColor = clamp(ink.rgb + ink.a * vec3(1.0, 0.15, 0.95), vec3(0.0), vec3(1.0));
                vec3 fire = mix(base, inkColor, clamp(max4(ink), 0.0, 1.0));

                vec3 color = fire * (0.12 * recovering + burning) + base * 0.03 * recovering;
                outColor = vec4(color, max(burning, recovering));
            }
        """.trimIndent()

        val STATE_UPDATE_FRAGMENT = """
            #version 450

            layout(set = 0, binding = 0) uniform sampler2D sourceTex;
            layout(set = 0, binding = 1) uniform sampler2D ignitionTex;

            layout(push_constant) uniform Push {
                vec4 refractoryDecay;
                float spreadStepPx;
                float spreadGain;
                float paintGain;
                float ignitionThreshold;
                float activeFireThreshold;
                float readyThreshold;
                float treeDensity;
                float treeScalePx;
                float treeJitter;
                int seed;
                float clearRequested;
                float spreadProbability;
                int frame;
            } push;

            layout(location = 0) out vec4 outColor;

            uint hash(uint x) {
                x ^= x >> 16u; x *= 0x7feb352du;
                x ^= x >> 15u; x *= 0x846ca68bu;
                x ^= x >> 16u;
                return x;
            }

            float random01(uvec2 p, uint salt) {
                return float(hash(p.x ^ hash(p.y ^ hash(uint(push.seed) ^ salt)))) / 4294967295.0;
            }

            // clumps on a coarse grid, roughened per texel by treeJitter
            bool isTree(ivec2 pixel) {
                uvec2 coarse = uvec2(floor(vec2(pixel) / max(push.treeScalePx, 1.0)));
                float mixed = mix(random01(coarse, 1u), random01(uvec2(pixel), 2u), clamp(push.treeJitter, 0.0, 1.0));
                return mixed <= push.treeDensity;
            }

            float max4(vec4 v) {
                return max(max(v.x, v.y), max(v.z, v.w));
            }

            // how strongly each kind is burning, 0 at activeFireThreshold, 1 at full
            vec4 burning(ivec2 p, ivec2 dims) {
                vec4 s = texelFetch(sourceTex, clamp(p, ivec2(0), dims - 1), 0);
                return clamp((s - vec4(push.activeFireThreshold)) / max(1.0 - push.activeFireThreshold, 0.0001), vec4(0.0), vec4(1.0));
            }

            // a burning neighbour's fire, if this frame's draw lets it across to this tree
            vec4 reaching(ivec2 pixel, ivec2 offset, ivec2 dims, uint salt) {
                uvec2 u = uvec2(pixel);
                float draw = random01(u, hash(uint(push.frame) * 8u + salt) | 4u);
                return draw < push.spreadProbability ? burning(pixel + offset, dims) : vec4(0.0);
            }

            // the kind with the strongest signal at or above the threshold, among those allowed; -1 for none
            int strongest(vec4 signal, bvec4 allowed) {
                int winner = -1;
                float best = -1.0;
                for (int k = 0; k < 4; ++k) {
                    if (allowed[k] && signal[k] >= push.ignitionThreshold && signal[k] > best) {
                        best = signal[k];
                        winner = k;
                    }
                }
                return winner;
            }

            vec4 caught(int kind) {
                vec4 v = vec4(0.0);
                v[kind] = 1.0;
                return v;
            }

            void main() {
                ivec2 dims = textureSize(sourceTex, 0);
                ivec2 pixel = clamp(ivec2(gl_FragCoord.xy), ivec2(0), dims - 1);

                if (push.clearRequested > 0.5 || !isTree(pixel)) {
                    outColor = vec4(0.0);
                    return;
                }

                vec4 previous = texelFetch(sourceTex, pixel, 0);
                vec4 decayed = max(previous - push.refractoryDecay, vec4(0.0));

                // drawn ignition takes any tree, burning or not
                int lit = strongest(texelFetch(ignitionTex, pixel, 0) * push.paintGain, bvec4(true));
                if (lit >= 0) {
                    outColor = caught(lit);
                    return;
                }

                if (max4(previous) >= push.activeFireThreshold) {
                    outColor = decayed;
                    return;
                }

                // each of the eight neighbours gets its own draw against spreadProbability
                int s = max(int(round(push.spreadStepPx)), 1);
                vec4 neighbours = max(max(max(reaching(pixel, ivec2(s, 0), dims, 0u), reaching(pixel, ivec2(-s, 0), dims, 1u)),
                                          max(reaching(pixel, ivec2(0, s), dims, 2u), reaching(pixel, ivec2(0, -s), dims, 3u))),
                                      0.7071 * max(max(reaching(pixel, ivec2(s, s), dims, 4u), reaching(pixel, ivec2(-s, s), dims, 5u)),
                                                   max(reaching(pixel, ivec2(s, -s), dims, 6u), reaching(pixel, ivec2(-s, -s), dims, 7u))));

                // a neighbour's fire takes a tree only if that kind has recovered there
                int spread = strongest(neighbours * push.spreadGain, lessThanEqual(decayed, vec4(push.readyThreshold)));
                outColor = spread >= 0 ? caught(spread) : decayed;
            }
        """.trimIndent()
    }
}