Compose shader glitch modifier

ShaderGlitchModifier.kt

package io.mightyseal.glitch

import android.graphics.RuntimeShader
import androidx.compose.animation.core.Animatable
import androidx.compose.animation.core.LinearEasing
import androidx.compose.animation.core.tween
import androidx.compose.runtime.Composable
import androidx.compose.runtime.LaunchedEffect
import androidx.compose.runtime.Stable
import androidx.compose.runtime.derivedStateOf
import androidx.compose.runtime.getValue
import androidx.compose.runtime.mutableIntStateOf
import androidx.compose.runtime.remember
import androidx.compose.runtime.setValue
import androidx.compose.ui.Modifier
import androidx.compose.ui.draw.clipToBounds
import androidx.compose.ui.graphics.graphicsLayer
import kotlin.math.roundToInt
import kotlin.random.Random

private const val DURATION = 500

@Stable
@Composable
fun Modifier.shaderGlitchEffect(
    key: Any? = null,
    slices: Int = 20,
): Modifier {
    val shader = remember {
        RuntimeShader(GLITCH_DEMO)
    }

    var step by remember { mutableIntStateOf(0) }
    val intensity by remember { derivedStateOf { step.toFloat() / 10f } }

    LaunchedEffect(key) {
        Animatable(10f)
            .animateTo(
                targetValue = 0f,
                animationSpec = tween(
                    durationMillis = DURATION,
                    easing = LinearEasing,
                ),
            ) {
                step = this.value.roundToInt()
            }
    }

    return clipToBounds()
        .graphicsLayer {
            if (intensity >= 0f) {
                shader.setIntUniform("slices", slices)
                shader.setFloatUniform("imageSize", this.size.width, this.size.height)
                shader.setFloatUniform("intensity", intensity)
                shader.setFloatUniform("realRandom", Random.nextFloat())

                renderEffect = shader.asEffect("image")
            }
        }
    }
}

private fun RuntimeShader.asEffect(uniformName: String): RenderEffect {
    return android.graphics.RenderEffect
        .createRuntimeShaderEffect(this, uniformName)
        .asComposeRenderEffect()
}

ShaderSource.kt

package io.mightyseal.glitch

import org.intellij.lang.annotations.Language

@Language("AGSL")
private val GLITCH_DEMO =
    """
    uniform shader image;
    uniform float2 imageSize; // Shader area size in pixels
    uniform float intensity;
    uniform int slices;
    uniform float realRandom;
    
    vec4 yellow = vec4(1.0, 1.0, 0.0, 1.0);
    vec4 red = vec4(1.0, 0.0, 0.0, 1.0);
    vec4 cyan = vec4(0.0, 1.0, 1.0, 1.0);
    
    // Simple random functions
    float random(float seed) {
        return fract(sin(seed) * 100000.0);
    }
    
    float random(float2 seed) {
        return fract(sin(dot(seed, float2(12.9898, 78.233))) * 43758.5453123);
    }

    // Determine how much this slice should be displaced
    float displace(float sliceY, float intensity) {
        float rnd = random(float2(sliceY, intensity));
        
        float shouldDisplace = 0.0;
        if (intensity < 0.5 && intensity > rnd * 0.4) {
            shouldDisplace = 0.0;
        } else {
            shouldDisplace = 1.0;
        }

        return (rnd - 0.5) * 40.0 * intensity * shouldDisplace;
    }
    
    
    float2 scale(float2 coord, float yMin, float yMax, float screenWidth, float intensity) {
        float rnd = random(float2(yMin, intensity));

        if (rnd < intensity) {
            float centerX = screenWidth * 0.5;
            
            float localX = coord.x - centerX;
            float scaleFactor = 1.0 + (intensity * rnd);
            localX /= scaleFactor;
            
            float scaledX = centerX + localX;
            
            return float2(scaledX, coord.y);
        }
        
        return coord;
    }

    half4 main(float2 fragCoord) {
        // Create horizontal slices
        float sliceHeight = imageSize.y / float(slices); // Height of each slice in pixels
        float sliceY = floor(fragCoord.y / sliceHeight) * sliceHeight; // Start coordinates for each slice

        float displace = displace(sliceY, intensity);
        float2 displaced = fragCoord;
        displaced.x += displace;


        float rnd = random(float2(intensity * realRandom, sliceY));
        if ((rnd * 2.5 + 0.5) < intensity) {
           if (realRandom > 0.67) {
               return yellow;
           } else if (realRandom > 0.33) {
               return red;
           } else {
               return cyan;
           }
        } else {
           float2 scaled = scale(displaced, sliceY, sliceY + sliceHeight, imageSize.y, intensity);
           return image.eval(scaled);
        }
    }
"""