1// NVIDIA FXAA 3.11
2// Original source code by TIMOTHY LOTTES
3// https://gist.github.com/kosua20/0c506b81b3812ac900048059d2383126
4//
5// Cleaned version - https://github.com/kosua20/Rendu/blob/master/resources/common/shaders/screens/fxaa.frag
6//
7// Tweaks by mrDIMAS - https://github.com/FyroxEngine/Fyrox/blob/master/src/renderer/shaders/fxaa_fs.glsl
8
9#import bevy_core_pipeline::fullscreen_vertex_shader::FullscreenVertexOutput
10
11@group(0) @binding(0) var screenTexture: texture_2d<f32>;
12@group(0) @binding(1) var samp: sampler;
13
14// Trims the algorithm from processing darks.
15#ifdef EDGE_THRESH_MIN_LOW
16    const EDGE_THRESHOLD_MIN: f32 = 0.0833;
17#endif
18
19#ifdef EDGE_THRESH_MIN_MEDIUM
20    const EDGE_THRESHOLD_MIN: f32 = 0.0625;
21#endif
22
23#ifdef EDGE_THRESH_MIN_HIGH
24    const EDGE_THRESHOLD_MIN: f32 = 0.0312;
25#endif
26
27#ifdef EDGE_THRESH_MIN_ULTRA
28    const EDGE_THRESHOLD_MIN: f32 = 0.0156;
29#endif
30
31#ifdef EDGE_THRESH_MIN_EXTREME
32    const EDGE_THRESHOLD_MIN: f32 = 0.0078;
33#endif
34
35// The minimum amount of local contrast required to apply algorithm.
36#ifdef EDGE_THRESH_LOW
37    const EDGE_THRESHOLD_MAX: f32 = 0.250;
38#endif
39
40#ifdef EDGE_THRESH_MEDIUM
41    const EDGE_THRESHOLD_MAX: f32 = 0.166;
42#endif
43
44#ifdef EDGE_THRESH_HIGH
45    const EDGE_THRESHOLD_MAX: f32 = 0.125;
46#endif
47
48#ifdef EDGE_THRESH_ULTRA
49    const EDGE_THRESHOLD_MAX: f32 = 0.063;
50#endif
51
52#ifdef EDGE_THRESH_EXTREME
53    const EDGE_THRESHOLD_MAX: f32 = 0.031;
54#endif
55
56const ITERATIONS: i32 = 12; //default is 12
57const SUBPIXEL_QUALITY: f32 = 0.75;
58// #define QUALITY(q) ((q) < 5 ? 1.0 : ((q) > 5 ? ((q) < 10 ? 2.0 : ((q) < 11 ? 4.0 : 8.0)) : 1.5))
59fn QUALITY(q: i32) -> f32 {
60    switch (q) {
61        //case 0, 1, 2, 3, 4: { return 1.0; }
62        default:              { return 1.0; }
63        case 5:               { return 1.5; }
64        case 6, 7, 8, 9:      { return 2.0; }
65        case 10:              { return 4.0; }
66        case 11:              { return 8.0; }
67    }
68}
69
70fn rgb2luma(rgb: vec3<f32>) -> f32 {
71    return sqrt(dot(rgb, vec3<f32>(0.299, 0.587, 0.114)));
72}
73
74// Performs FXAA post-process anti-aliasing as described in the Nvidia FXAA white paper and the associated shader code.
75@fragment
76fn fragment(in: FullscreenVertexOutput) -> @location(0) vec4<f32> {
77    let resolution = vec2<f32>(textureDimensions(screenTexture));
78    let inverseScreenSize = 1.0 / resolution.xy;
79    let texCoord = in.position.xy * inverseScreenSize;
80
81    let centerSample = textureSampleLevel(screenTexture, samp, texCoord, 0.0);
82    let colorCenter = centerSample.rgb;
83
84    // Luma at the current fragment
85    let lumaCenter = rgb2luma(colorCenter);
86
87    // Luma at the four direct neighbors of the current fragment.
88    let lumaDown = rgb2luma(textureSampleLevel(screenTexture, samp, texCoord, 0.0, vec2<i32>(0, -1)).rgb);
89    let lumaUp = rgb2luma(textureSampleLevel(screenTexture, samp, texCoord, 0.0, vec2<i32>(0, 1)).rgb);
90    let lumaLeft = rgb2luma(textureSampleLevel(screenTexture, samp, texCoord, 0.0, vec2<i32>(-1, 0)).rgb);
91    let lumaRight = rgb2luma(textureSampleLevel(screenTexture, samp, texCoord, 0.0, vec2<i32>(1, 0)).rgb);
92
93    // Find the maximum and minimum luma around the current fragment.
94    let lumaMin = min(lumaCenter, min(min(lumaDown, lumaUp), min(lumaLeft, lumaRight)));
95    let lumaMax = max(lumaCenter, max(max(lumaDown, lumaUp), max(lumaLeft, lumaRight)));
96
97    // Compute the delta.
98    let lumaRange = lumaMax - lumaMin;
99
100    // If the luma variation is lower that a threshold (or if we are in a really dark area), we are not on an edge, don't perform any AA.
101    if (lumaRange < max(EDGE_THRESHOLD_MIN, lumaMax * EDGE_THRESHOLD_MAX)) {
102        return centerSample;
103    }
104
105    // Query the 4 remaining corners lumas.
106    let lumaDownLeft  = rgb2luma(textureSampleLevel(screenTexture, samp, texCoord, 0.0, vec2<i32>(-1, -1)).rgb);
107    let lumaUpRight   = rgb2luma(textureSampleLevel(screenTexture, samp, texCoord, 0.0, vec2<i32>(1, 1)).rgb);
108    let lumaUpLeft    = rgb2luma(textureSampleLevel(screenTexture, samp, texCoord, 0.0, vec2<i32>(-1, 1)).rgb);
109    let lumaDownRight = rgb2luma(textureSampleLevel(screenTexture, samp, texCoord, 0.0, vec2<i32>(1, -1)).rgb);
110
111    // Combine the four edges lumas (using intermediary variables for future computations with the same values).
112    let lumaDownUp = lumaDown + lumaUp;
113    let lumaLeftRight = lumaLeft + lumaRight;
114
115    // Same for corners
116    let lumaLeftCorners = lumaDownLeft + lumaUpLeft;
117    let lumaDownCorners = lumaDownLeft + lumaDownRight;
118    let lumaRightCorners = lumaDownRight + lumaUpRight;
119    let lumaUpCorners = lumaUpRight + lumaUpLeft;
120
121    // Compute an estimation of the gradient along the horizontal and vertical axis.
122    let edgeHorizontal = abs(-2.0 * lumaLeft   + lumaLeftCorners)  + 
123                         abs(-2.0 * lumaCenter + lumaDownUp) * 2.0 + 
124                         abs(-2.0 * lumaRight  + lumaRightCorners);
125
126    let edgeVertical =   abs(-2.0 * lumaUp     + lumaUpCorners)       + 
127                         abs(-2.0 * lumaCenter + lumaLeftRight) * 2.0 + 
128                         abs(-2.0 * lumaDown   + lumaDownCorners);
129
130    // Is the local edge horizontal or vertical ?
131    let isHorizontal = (edgeHorizontal >= edgeVertical);
132
133    // Choose the step size (one pixel) accordingly.
134    var stepLength = select(inverseScreenSize.x, inverseScreenSize.y, isHorizontal);
135
136    // Select the two neighboring texels lumas in the opposite direction to the local edge.
137    var luma1 = select(lumaLeft, lumaDown, isHorizontal);
138    var luma2 = select(lumaRight, lumaUp, isHorizontal);
139
140    // Compute gradients in this direction.
141    let gradient1 = luma1 - lumaCenter;
142    let gradient2 = luma2 - lumaCenter;
143
144    // Which direction is the steepest ?
145    let is1Steepest = abs(gradient1) >= abs(gradient2);
146
147    // Gradient in the corresponding direction, normalized.
148    let gradientScaled = 0.25 * max(abs(gradient1), abs(gradient2));
149
150    // Average luma in the correct direction.
151    var lumaLocalAverage = 0.0;
152    if (is1Steepest) {
153        // Switch the direction
154        stepLength = -stepLength;
155        lumaLocalAverage = 0.5 * (luma1 + lumaCenter);
156    } else {
157        lumaLocalAverage = 0.5 * (luma2 + lumaCenter);
158    }
159
160    // Shift UV in the correct direction by half a pixel.
161    // Compute offset (for each iteration step) in the right direction.
162    var currentUv = texCoord;
163    var offset = vec2<f32>(0.0, 0.0);
164    if (isHorizontal) {
165        currentUv.y = currentUv.y + stepLength * 0.5;
166        offset.x = inverseScreenSize.x;
167    } else {
168        currentUv.x = currentUv.x + stepLength * 0.5;
169        offset.y = inverseScreenSize.y;
170    }
171
172    // Compute UVs to explore on each side of the edge, orthogonally. The QUALITY allows us to step faster.
173    var uv1 = currentUv - offset; // * QUALITY(0); // (quality 0 is 1.0)
174    var uv2 = currentUv + offset; // * QUALITY(0); // (quality 0 is 1.0)
175
176    // Read the lumas at both current extremities of the exploration segment, and compute the delta wrt to the local average luma.
177    var lumaEnd1 = rgb2luma(textureSampleLevel(screenTexture, samp, uv1, 0.0).rgb);
178    var lumaEnd2 = rgb2luma(textureSampleLevel(screenTexture, samp, uv2, 0.0).rgb);
179    lumaEnd1 = lumaEnd1 - lumaLocalAverage;
180    lumaEnd2 = lumaEnd2 - lumaLocalAverage;
181
182    // If the luma deltas at the current extremities is larger than the local gradient, we have reached the side of the edge.
183    var reached1 = abs(lumaEnd1) >= gradientScaled;
184    var reached2 = abs(lumaEnd2) >= gradientScaled;
185    var reachedBoth = reached1 && reached2;
186
187    // If the side is not reached, we continue to explore in this direction.
188    uv1 = select(uv1 - offset, uv1, reached1); // * QUALITY(1); // (quality 1 is 1.0)
189    uv2 = select(uv2 - offset, uv2, reached2); // * QUALITY(1); // (quality 1 is 1.0)
190
191    // If both sides have not been reached, continue to explore.
192    if (!reachedBoth) {
193        for (var i: i32 = 2; i < ITERATIONS; i = i + 1) {
194            // If needed, read luma in 1st direction, compute delta.
195            if (!reached1) { 
196                lumaEnd1 = rgb2luma(textureSampleLevel(screenTexture, samp, uv1, 0.0).rgb);
197                lumaEnd1 = lumaEnd1 - lumaLocalAverage;
198            }
199            // If needed, read luma in opposite direction, compute delta.
200            if (!reached2) { 
201                lumaEnd2 = rgb2luma(textureSampleLevel(screenTexture, samp, uv2, 0.0).rgb);
202                lumaEnd2 = lumaEnd2 - lumaLocalAverage;
203            }
204            // If the luma deltas at the current extremities is larger than the local gradient, we have reached the side of the edge.
205            reached1 = abs(lumaEnd1) >= gradientScaled;
206            reached2 = abs(lumaEnd2) >= gradientScaled;
207            reachedBoth = reached1 && reached2;
208
209            // If the side is not reached, we continue to explore in this direction, with a variable quality.
210            if (!reached1) {
211                uv1 = uv1 - offset * QUALITY(i);
212            }
213            if (!reached2) {
214                uv2 = uv2 + offset * QUALITY(i);
215            }
216
217            // If both sides have been reached, stop the exploration.
218            if (reachedBoth) { 
219                break; 
220            }
221        }
222    }
223
224    // Compute the distances to each side edge of the edge (!).
225    var distance1 = select(texCoord.y - uv1.y, texCoord.x - uv1.x, isHorizontal);
226    var distance2 = select(uv2.y - texCoord.y, uv2.x - texCoord.x, isHorizontal);
227
228    // In which direction is the side of the edge closer ?
229    let isDirection1 = distance1 < distance2;
230    let distanceFinal = min(distance1, distance2);
231
232    // Thickness of the edge.
233    let edgeThickness = (distance1 + distance2);
234
235    // Is the luma at center smaller than the local average ?
236    let isLumaCenterSmaller = lumaCenter < lumaLocalAverage;
237
238    // If the luma at center is smaller than at its neighbor, the delta luma at each end should be positive (same variation).
239    let correctVariation1 = (lumaEnd1 < 0.0) != isLumaCenterSmaller;
240    let correctVariation2 = (lumaEnd2 < 0.0) != isLumaCenterSmaller;
241
242    // Only keep the result in the direction of the closer side of the edge.
243    var correctVariation = select(correctVariation2, correctVariation1, isDirection1);
244
245    // UV offset: read in the direction of the closest side of the edge.
246    let pixelOffset = - distanceFinal / edgeThickness + 0.5;
247
248    // If the luma variation is incorrect, do not offset.
249    var finalOffset = select(0.0, pixelOffset, correctVariation);
250
251    // Sub-pixel shifting
252    // Full weighted average of the luma over the 3x3 neighborhood.
253    let lumaAverage = (1.0 / 12.0) * (2.0 * (lumaDownUp + lumaLeftRight) + lumaLeftCorners + lumaRightCorners);
254    // Ratio of the delta between the global average and the center luma, over the luma range in the 3x3 neighborhood.
255    let subPixelOffset1 = clamp(abs(lumaAverage - lumaCenter) / lumaRange, 0.0, 1.0);
256    let subPixelOffset2 = (-2.0 * subPixelOffset1 + 3.0) * subPixelOffset1 * subPixelOffset1;
257    // Compute a sub-pixel offset based on this delta.
258    let subPixelOffsetFinal = subPixelOffset2 * subPixelOffset2 * SUBPIXEL_QUALITY;
259
260    // Pick the biggest of the two offsets.
261    finalOffset = max(finalOffset, subPixelOffsetFinal);
262
263    // Compute the final UV coordinates.
264    var finalUv = texCoord;
265    if (isHorizontal) {
266        finalUv.y = finalUv.y + finalOffset * stepLength;
267    } else {
268        finalUv.x = finalUv.x + finalOffset * stepLength;
269    }
270
271    // Read the color at the new UV coordinates, and use it.
272    var finalColor = textureSampleLevel(screenTexture, samp, finalUv, 0.0).rgb;
273    return vec4<f32>(finalColor, centerSample.a);
274}