Jetpack Compose Ripple shader

Shader.kt

// Credits https://kotlinlang.slack.com/archives/CJLTWPH7S/p1718623704772239
import android.graphics.RenderEffect
import android.graphics.RuntimeShader
import android.os.Build
import androidx.annotation.RequiresApi
import androidx.compose.animation.core.Animatable
import androidx.compose.animation.core.LinearEasing
import androidx.compose.animation.core.tween
import androidx.compose.runtime.LaunchedEffect
import androidx.compose.runtime.getValue
import androidx.compose.runtime.mutableStateOf
import androidx.compose.runtime.remember
import androidx.compose.runtime.setValue
import androidx.compose.ui.Modifier
import androidx.compose.ui.composed
import androidx.compose.ui.geometry.Offset
import androidx.compose.ui.graphics.asComposeRenderEffect
import androidx.compose.ui.graphics.graphicsLayer
import androidx.compose.ui.layout.onSizeChanged
import androidx.compose.ui.platform.LocalDensity
import androidx.compose.ui.unit.dp
import androidx.compose.ui.util.fastRoundToInt
import kotlin.math.sqrt
import org.intellij.lang.annotations.Language

@Language(value = "AGSL")
val shaderCode = """
uniform shader image;

uniform float2 origin;
uniform float elapsedTime;
uniform float amplitude;
uniform float frequency;
uniform float decay;
uniform float speed;
uniform float radius;
uniform float tintRatio;

float2 lerp(float2 a, float2 b, float t) {
    return a + t * (b - a);
}

// Note: we use easing to prevent deformation in the center

// Function to calculate the cubic Bezier curve value at t
vec2 cubicBezier(vec2 P0, vec2 P1, vec2 P2, vec2 P3, float t) {
    float u = 1.0 - t;
    float tt = t * t;
    float uu = u * u;
    float uuu = uu * u;
    float ttt = tt * t;

    vec2 p = uuu * P0;       // (1 - t)^3 * P0
    p += 3.0 * uu * t * P1;  // 3 * (1 - t)^2 * t * P1
    p += 3.0 * u * tt * P2;  // 3 * (1 - t) * t^2 * P2
    p += ttt * P3;           // t^3 * P3

    return p;
}

// Function to find t for a given x using the Newton-Raphson method
float findTForX(float x, vec2 P0, vec2 P1, vec2 P2, vec2 P3) {
    float t = x; // Initial guess
    for (int i = 0; i < 5; i++) { // Iterate a few times
        vec2 bezierPoint = cubicBezier(P0, P1, P2, P3, t);
        float x_t = bezierPoint.x;
        float dx_dt = -3.0 * (1.0 - t) * (1.0 - t) * P0.x + 3.0 * (1.0 - 2.0 * t) * (1.0 - t) * P1.x + 3.0 * t * (2.0 - 3.0 * t) * P2.x + 3.0 * t * t * P3.x;

        t -= (x_t - x) / dx_dt;
        t = clamp(t, 0.0, 1.0); // Keep t within [0, 1]
    }
    return t;
}

// Function to get the y value for a given x using the cubic Bezier curve
float getYForX(float x, vec2 P1, vec2 P2) {
    vec2 P0 = vec2(0.0, 0.0);  // Start point
    vec2 P3 = vec2(1.0, 1.0);  // End point

    float t = findTForX(x, P0, P1, P2, P3);
    vec2 bezierPoint = cubicBezier(P0, P1, P2, P3, t);

    return bezierPoint.y;
}

float transformEaseEmphasizedAccelerate(float x) {
    float2 P1 = float2(0.05, 0.7);
    float2 P2 = float2(0.1, 1.0);
    return getYForX(x, P1, P2);
}

half4 main(float2 coord) {
    float distance = length(coord - origin);
    float delay = distance / speed;
    float adjustedTime = max(0.0, elapsedTime - delay);

    float rippleAmount = amplitude * sin(frequency * adjustedTime) * exp(-decay * adjustedTime);
    float2 n = normalize(coord - origin);
    float2 newPosition = coord + rippleAmount * n;
    newPosition = lerp(coord, newPosition, transformEaseEmphasizedAccelerate(distance / radius));

    half4 color = image.eval(newPosition);
    if (amplitude != 0.0) {
        color.rgb += tintRatio * (rippleAmount / amplitude) * color.a;
    }
    return color;
}	
""".trimIndent()

@RequiresApi(Build.VERSION_CODES.TIRAMISU)
fun Modifier.shader() = this.composed {
	var shader by remember { mutableStateOf<RippleShader?>(null) }
	val density = LocalDensity.current
	LaunchedEffect(Unit) {
		while (true) {
			shader?.fadeIn()
			shader?.fadeOut()
		}
	}
	Modifier
		.graphicsLayer {
			renderEffect = shader?.renderEffect
			clip = true
		}
		.onSizeChanged { size ->
			shader = RippleShader(
				origin = Offset(size.width / 2f, size.height / 1.25f),
				radius = sqrt(size.width.toFloat() * size.width + size.height * size.height) / 2f,
				amplitude = with(density) { (-16).dp.toPx() },
				frequency = 10f,
				decay = 10f,
				relativeSpeed = 6f,
				tintRatio = 0.3f,
				startTimeMillis = 0,
				fadeInDurationMillis = 0,
				fadeOutDurationMillis = 500
			)
		}
}

@RequiresApi(33)
class RippleShader(
	val origin: Offset,
	val radius: Float,
	val amplitude: Float,
	val frequency: Float,
	val decay: Float,
	val relativeSpeed: Float,
	val tintRatio: Float,
	val startTimeMillis: Int,
	val fadeInDurationMillis: Int,
	val fadeOutDurationMillis: Int
) {
	private val elapsedTime = Animatable(startTimeMillis / 1000f, visibilityThreshold = 1f)

	private val runtimeShader = RuntimeShader(shaderCode).apply {
		setFloatUniform("origin", origin.x, origin.y)
		setFloatUniform("elapsedTime", 0f)
		setFloatUniform("amplitude", amplitude)
		setFloatUniform("frequency", frequency)
		setFloatUniform("decay", decay)
		setFloatUniform("speed", radius * relativeSpeed)
		setFloatUniform("radius", radius)
		setFloatUniform("tintRatio", tintRatio)
	}

	var renderEffect by mutableStateOf(
		RenderEffect
			.createRuntimeShaderEffect(runtimeShader, "image")
			.asComposeRenderEffect()
	)
		private set

	suspend fun fadeIn() {
		elapsedTime.animateTo(
			(startTimeMillis + fadeInDurationMillis) / 1000f,
			tween(
				startTimeMillis + fadeInDurationMillis - (elapsedTime.value * 1000).fastRoundToInt(),
				0,
				LinearEasing
			)
		) {
			runtimeShader.setFloatUniform("elapsedTime", value)
			renderEffect = RenderEffect
				.createRuntimeShaderEffect(runtimeShader, "image")
				.asComposeRenderEffect()
		}
	}

	suspend fun fadeOut() {
		elapsedTime.animateTo(
			(startTimeMillis + fadeInDurationMillis + fadeOutDurationMillis) / 1000f,
			tween(fadeOutDurationMillis, 0, LinearEasing)
		) {
			runtimeShader.setFloatUniform("elapsedTime", value)
			renderEffect = RenderEffect
				.createRuntimeShaderEffect(runtimeShader, "image")
				.asComposeRenderEffect()
		}
		elapsedTime.snapTo(0f)
	}
}