Reland cache snow flakes and reduce snow layers

    highp vec2 cellUv;

};

const mat2 rot45 = mat2(

    0.7071067812, 0.7071067812, // First column.

    -0.7071067812, 0.7071067812 // second column.

);

const vec2 snowFlakeShape = vec2(0.28, 0.26);

// decreasedFactor should match minDescreasedFactor * 2, minDescreasedFactor is defined in snow_flake_samples.agsl

const float decreasedFactor = 1.0 / 0.28;

uniform half intensity;

const float farthestSnowLayerWiggleSpeed = 2.18;

const float closestSnowLayerWiggleSpeed = 0.9;

uniform half snowFlakeSamplesSize;

/**

 * Generates snow flakes.

    in float minLayerIndex,

    in float maxLayerIndex

) {

    // Normalize the layer index. 0 is closest, 1 is farthest.

    half normalizedLayerIndex = map(layerIndex, minLayerIndex, maxLayerIndex, 0, 1);

    /* Grid. */

    // Increase the last number to make each layer more separate from the previous one.

    float depth = 0.65 + layerIndex * 0.555;

    float depth = 0.65 + layerIndex * 0.755;

    float speedAdj = 1. + layerIndex * 0.225;

    float layerR = idGenerator(layerIndex);

    snowGridSize *= depth;

    time += layerR * 58.3;

    // Number of rows and columns (each one is a cell, a drop).

    float cellAspectRatio = snowGridSize.x / snowGridSize.y;

    // Aspect ratio impacts visible cells.

    uv.y /= screenAspectRatio;

    // Skew uv.x so it goes to left or right

    // Have time affect the position of each column as well.

    gridUv.y += columnId * 2.6 + time * 0.19 * (1 - columnId);

    /* Cell. */

    // Get the cell ID based on the grid position. Value from 0 to 1.

    float cellId = idGenerator(floor(gridUv));

    // For each cell, we set the internal UV from -0.5 (left, bottom) to 0.5 (right, top).

    // Calclulate the grid this pixel belonging to, and also the offset in the cell.

    vec2 gridIdx = floor(gridUv);

    vec2 cellUv = fract(gridUv) - 0.5;

    cellUv.y *= -1.;

   /*

    * Disable snow flakes with some probabilty. This is done by 1) assigning a random intensity

    * value to the cell 2) then compare it with the given intensity.

    */

    half cellIntensity = idGenerator(floor(vec2(cellId * 856.16, 272.2)));

    if (cellIntensity < 1. - intensity) {

        // Remove snow flakes by seeting flake mask to 0.

    // The bigger the decreasedFactor, the smaller the snow flake.

    vec2 snowFlakePos = vec2(cellUv.x, cellUv.y * (snowGridSize.x / snowGridSize.y - 1.0 / snowGridSize.y) + uv.y - 0.5 / screenAspectRatio) * decreasedFactor;

    if (abs(snowFlakePos.y) > 0.5 || abs(snowFlakePos.x) > 0.5 ) {

        return Snow(/* flakeMask= */ 0, cellUv);

    }

    vec4 color = snowFlakeSamples.eval(snowFlakeSamplesSize * (gridIdx - 0.5 + snowFlakePos));

    /* Cell-id-based variations. */

    // 0 = snow flake invisible, 1 = snow flake visible.

    float baseMask = color.r;

    half cellIntensity = color.g;

    float cellId = color.b;

    if (cellIntensity <= 1. - intensity) {

        // Remove snow flakes by seting flake mask to 0.

        return Snow(/* flakeMask= */ 0, cellUv);

    }

    float visibilityFactor = smoothstep(

        cellIntensity,

        max(cellIntensity - (0.02 + 0.18 * intensity), 0.0),

        1 - intensity);

    // Adjust the size of each snow flake (higher is smaller) based on cell ID.

    float decreaseFactor = 2.0 + map(cellId, 0., 1., -0.1, 2.8) + 5. * (1 - visibilityFactor);

    // Adjust the opacity of the particle based on the cell id and distance from the camera.

    float farLayerFadeOut = map(normalizedLayerIndex, 0.7, 1, 1, 0.4);

    float closeLayerFadeOut = map(normalizedLayerIndex, 0, 0.2, 0.6, 1);

    float opacityVariation =

        (1. - 0.9 * cellId)*

        visibilityFactor *

        closeLayerFadeOut *

        farLayerFadeOut;

    /* Cell snow flake. */

    // Calculate snow flake.

    vec2 snowFlakeShape = vec2(0.28, 0.26);

    vec2 snowFlakePos = vec2(cellUv.x, cellUv.y * cellAspectRatio);

    snowFlakePos -= vec2(

            0.,

            (uv.y - 0.5 / screenAspectRatio)  - cellUv.y / snowGridSize.y

        ) * screenAspectRatio;

    snowFlakePos *= snowFlakeShape * decreaseFactor;

    vec2 snowFlakeShapeVariation = vec2(0.055) * // max variation

        vec2((cellId * 2. - 1.), // random A based on cell ID

            (fract((cellId + 0.03521) * 34.21) * 2. - 1.)); // random B based on cell ID

    vec2 snowFlakePosR = 1.016 * abs(rot45 * (snowFlakePos + snowFlakeShapeVariation));

    snowFlakePos = abs(snowFlakePos);

    // Create the snowFlake mask.

    float flakeMask = smoothstep(

        0.3,

        0.200 - 0.3 * opacityVariation,

        snowFlakePos.x + snowFlakePos.y + snowFlakePosR.x + snowFlakePosR.y

    ) * opacityVariation;

        visibilityFactor * farLayerFadeOut * closeLayerFadeOut;

   float flakeMask = baseMask * opacityVariation;

    return Snow(flakeMask, cellUv);

}

uniform float time;

uniform float screenAspectRatio;

uniform float2 screenSize;

uniform float cellAspectRatio;

uniform mat3 transformMatrixBitmap;

uniform mat3 transformMatrixWeather;

uniform shader snowFlakeSamples;

#include "shaders/constants.agsl"

#include "shaders/utils.agsl"

// Snow tint.

const vec4 snowColor = vec4(1., 1., 1., 0.95);

// Background tint

const vec4 bgdTint = vec4(0.8, 0.8, 0.8, 0.07);

// Indices of the different snow layers.

const float farthestSnowLayerIndex = 6;

const float farthestSnowLayerIndex = 4;

const float midSnowLayerIndex = 2;

const float closestSnowLayerIndex = 0;

/*

 * Copyright (C) 2025 The Android Open Source Project

 *

 * Licensed under the Apache License, Version 2.0 (the "License");

 * you may not use this file except in compliance with the License.

 * You may obtain a copy of the License at

 *

 *      http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

uniform float2 canvasSize;

uniform float snowFlakeSamplesSize;

#include "shaders/constants.agsl"

#include "shaders/utils.agsl"

const vec2 snowFlakeShape = vec2(0.28, 0.26);

// Used in generate snow flake samples, and sampling it afterwards, to make sure snow flakes

// will not go beyond bounding box

const float minDecreaseFactor = 0.5 / 0.28;

const float layerIndex = 0;

const mat2 rot45 = mat2(

    0.7071067812, 0.7071067812, // First column.

    -0.7071067812, 0.7071067812 // second column.

);

/**

 * This shader generates snow flake samples per cell. It stores the flake mask in the red channel,

 * and pre-calculated `cellIntensity` and `cellId` in the green and blue channels for optimized access.

 */

vec4 main(float2 fragCoord) {

    // Calculate uv for snow based on transformed coordinates

    float2 uv = fragCoord / canvasSize;

    float layerR = idGenerator(layerIndex);

    // Number of rows and columns (each one is a cell, a snowflake).

    vec2 gridSize = floor(canvasSize / snowFlakeSamplesSize);

    float cellAspectRatio = gridSize.x / gridSize.y;

    // Aspect ratio impacts visible cells.

    vec2 gridUv = uv * gridSize;

    /* Cell. */

    // Get the cell ID based on the grid position. Value from 0 to 1.

    float cellId = idGenerator(floor(gridUv));

    // For each cell, we set the internal UV from -0.5 (left, bottom) to 0.5 (right, top).

    vec2 cellUv = fract(gridUv) - 0.5;

    cellUv.y *= -1.;

   /*

    * Disable snow flakes with some probabilty. This is done by 1) assigning a random intensity

    * value to the cell 2) then compare it with the given intensity.

    */

    half cellIntensity = idGenerator(floor(vec2(cellId * 856.16, 272.2)));

    /* Cell snow flake. */

    // Calculate snow flake.

    // With decreaseFactor <= minSnowShapeScale, we can make sure snow flakes not going out its

    // snowFlakeSamplesSize * snowFlakeSamplesSize bounding box

    float decreaseFactor = clamp(2.0 + map(cellId, 0., 1., 1., 1 + 5. * (1 - cellIntensity)),

        minDecreaseFactor, 4);

    // snowFlake center should be (0,0) in the cell when generating snowflake samples

    vec2 snowFlakePos = vec2(cellUv.x, cellUv.y);

    snowFlakePos *= snowFlakeShape * decreaseFactor;

    vec2 snowFlakeShapeVariation = vec2(0.055) * // max variation

        vec2((cellId * 2. - 1.), // random A based on cell ID

        (fract((cellId + 0.03521) * 34.21) * 2. - 1.)); // random B based on cell ID

    vec2 snowFlakePosR = 1.016 * abs(rot45 * (snowFlakePos + snowFlakeShapeVariation));

    snowFlakePos = abs(snowFlakePos);

    // Create the snowFlake mask.

    float baseMask = 1 - clamp(snowFlakePos.x + snowFlakePos.y + snowFlakePosR.x + snowFlakePosR.y, 0, 1);

    return vec4(baseMask, cellIntensity, cellId , 1);

}

    // Convert back to Cartesian coordinates (x, y)

    return transformedPoint.xy / transformedPoint.z;

}

float normalizeValue(float x, float minVal, float maxVal) {

    return (x - minVal) / (maxVal - minVal);

}

import android.graphics.RenderEffect

import android.graphics.RuntimeShader

import android.graphics.Shader

import android.hardware.HardwareBuffer

import android.util.SizeF

import androidx.core.graphics.createBitmap

import com.google.android.wallpaper.weathereffects.graphics.FrameBuffer

    private val accumulationFrameBufferPaint =

        Paint().also { it.shader = snowConfig.accumulatedSnowResultShader }

    private val snowFlakeSamplesBuffer: FrameBuffer =

        if (

            HardwareBuffer.isSupported(

                SNOW_FLAKE_SAMPLES_BUFFER_WIDTH,

                SNOW_FLAKE_SAMPLES_BUFFER_HEIGHT,

                HardwareBuffer.RGB_888,

                1,

                HardwareBuffer.USAGE_GPU_SAMPLED_IMAGE or HardwareBuffer.USAGE_GPU_COLOR_OUTPUT,

            )

        ) {

            FrameBuffer(

                SNOW_FLAKE_SAMPLES_BUFFER_WIDTH,

                SNOW_FLAKE_SAMPLES_BUFFER_HEIGHT,

                HardwareBuffer.RGB_888,

            )

        } else {

            FrameBuffer(SNOW_FLAKE_SAMPLES_BUFFER_WIDTH, SNOW_FLAKE_SAMPLES_BUFFER_HEIGHT)

        }

    private val snowFlakeSamplesPaint = Paint().also { it.shader = snowConfig.snowFlakeSamples }

    init {

        outlineFrameBuffer.setRenderEffect(

            RenderEffect.createBlurEffect(

                Shader.TileMode.CLAMP,

            )

        )

        updateTextureUniforms()

        adjustCropping(surfaceSize)

        prepareColorGrading()

        // Generate accumulated snow at the end after we updated all the uniforms.

        generateAccumulatedSnow()

        generateSnowFlakeSamples()

    }

    override fun update(deltaMillis: Long, frameTimeNanos: Long) {

        elapsedTime += snowSpeed * TimeUtils.millisToSeconds(deltaMillis)

        snowConfig.shader.setFloatUniform("time", elapsedTime)

        snowConfig.colorGradingShader.setInputShader("texture", snowConfig.shader)

    }

        super.release()

        outlineFrameBuffer.close()

        accumulationFrameBuffer.close()

        snowFlakeSamplesBuffer.close()

    }

    override fun setIntensity(intensity: Float) {

            "noise",

            BitmapShader(snowConfig.noiseTexture, Shader.TileMode.REPEAT, Shader.TileMode.REPEAT),

        )

        snowConfig.shader.setFloatUniform("snowFlakeSamplesSize", SNOW_FLAKE_SAMPLES_SIZE.toFloat())

        snowConfig.snowFlakeSamples.setFloatUniform(

            "snowFlakeSamplesSize",

            SNOW_FLAKE_SAMPLES_SIZE.toFloat(),

        )

    }

    private fun prepareColorGrading() {

        outlineFrameBuffer.endDrawing()

        outlineFrameBuffer.tryObtainingImage(

            ::generateAccumulatedSnowWithBlurredOutline,

            this::generateAccumulatedSnowWithBlurredOutline,

            mainExecutor,

        )

    }

    override fun updateGridSize(newSurfaceSize: SizeF) {

        val gridSize = GraphicsUtils.computeDefaultGridSize(newSurfaceSize, snowConfig.pixelDensity)

        snowConfig.shader.setFloatUniform("gridSize", 7 * gridSize, 2f * gridSize)

        val gridSizeColumns = 7f * gridSize

        val gridSizeRows = 2f * gridSize

        snowConfig.shader.setFloatUniform("gridSize", gridSizeColumns, gridSizeRows)

        snowConfig.shader.setFloatUniform("cellAspectRatio", gridSizeColumns / gridSizeRows)

    }

    /**

     * Generates an offscreen bitmap containing pre-rendered snow flake patterns and properties.

     *

     * This bitmap serves as a lookup table for the main snow_effect shader, reducing per-frame

     * calculations.

     */

    private fun generateSnowFlakeSamples() {

        val renderingCanvas = snowFlakeSamplesBuffer.beginDrawing()

        snowConfig.snowFlakeSamples.setFloatUniform(

            "canvasSize",

            SNOW_FLAKE_SAMPLES_BUFFER_WIDTH.toFloat(),

            SNOW_FLAKE_SAMPLES_BUFFER_HEIGHT.toFloat(),

        )

        renderingCanvas.drawPaint(snowFlakeSamplesPaint)

        snowFlakeSamplesBuffer.endDrawing()

        snowFlakeSamplesBuffer.tryObtainingImage(

            { snowFlakeSamples ->

                snowConfig.shader.setInputShader(

                    "snowFlakeSamples",

                    BitmapShader(snowFlakeSamples, Shader.TileMode.MIRROR, Shader.TileMode.MIRROR),

                )

            },

            mainExecutor,

        )

    }

    companion object {

        // To prevent displaying outdated accumulation, we use a tiny blank bitmap to temporarily

        // clear the rendering area before the new texture is ready.

        private val blankBitmap = createBitmap(1, 1)

        // The `snow_flakes_samples` shader pre-generates diverse snow flake properties

        // (shape mask, intensity, etc.) in a bitmap, reducing per-frame calculations. A higher

        // column count provides more x-based visual variations.

        // The following values balance the visual benefits with memory and shader sampling costs.

        // Number of columns; increases x-based variation.

        const val SNOW_FLAKE_SAMPLES_COLUMN_COUNT = 14

        const val SNOW_FLAKE_SAMPLES_ROW_COUNT = 4

        // Side length of each flake's square bounding box.

        const val SNOW_FLAKE_SAMPLES_SIZE = 100

        const val SNOW_FLAKE_SAMPLES_BUFFER_WIDTH =

            SNOW_FLAKE_SAMPLES_COLUMN_COUNT * SNOW_FLAKE_SAMPLES_SIZE

        const val SNOW_FLAKE_SAMPLES_BUFFER_HEIGHT =

            SNOW_FLAKE_SAMPLES_ROW_COUNT * SNOW_FLAKE_SAMPLES_SIZE

    }

}