ShaderGradient.metal

#include <metal_stdlib>
using namespace metal;

typedef enum ShaderGradientInputIndex
{
    ShaderGradientInputIndexVertices = 0,
    ShaderGradientInputIndexModelConstants,
    ShaderGradientInputIndexSceneConstants
} ShaderGradientInputIndex;

struct VertexIn {
    float4 position;
    float4 color;
    float2 index;
};

struct ModelConstants {
    float4x4 modelMatrix;
    int2 vertexCount;
};

struct PlaneModifiers {
    float noiseDencity;
    float noiseStrength;
    float outMultiplier;
    float inMultiplier;
    float normalX;
    float normalY;
    float normalZ;
    float t;
};

struct Coloring {
    float4 topColor;
    float4 bottomColor;
    float topBorder;
    float bottomBorder;
};

struct EffectsConstants {
    float radius;
    float rotateR;
    float rotateG;
    float rotateB;
    float scatter;
    float width;
    float height;
};

struct SceneConstants {
    float4x4 viewMatrix;
    float4x4 projection;
    struct PlaneModifiers planeModifiers;
    struct Coloring coloring;
    struct EffectsConstants effects;
};

// MARK: - Vertex shader

float3 mod289(float3 x)
{
    return x - floor(x * (1.0 / 289.0)) * 289.0;
}

float4 mod289(float4 x)
{
    return x - floor(x * (1.0 / 289.0)) * 289.0;
}

float4 permute(float4 x)
{
    return mod289(((x*34.0)+1.0)*x);
}

float4 taylorInvSqrt(float4 r)
{
    return 1.79284291400159 - 0.85373472095314 * r;
}

float3 fade(float3 t) {
    return t*t*t*(t*(t*6.0-15.0)+10.0);
}

float cnoise(float3 P) {
    float3 Pi0 = floor(P); // Integer part for indexing
    float3 Pi1 = Pi0 + float3(1.0); // Integer part + 1
    Pi0 = mod289(Pi0);
    Pi1 = mod289(Pi1);
    float3 Pf0 = fract(P); // Fractional part for interpolation
    float3 Pf1 = Pf0 - float3(1.0); // Fractional part - 1.0
    float4 ix = float4(Pi0.x, Pi1.x, Pi0.x, Pi1.x);
    float4 iy = float4(Pi0.yy, Pi1.yy);
    float4 iz0 = Pi0.zzzz;
    float4 iz1 = Pi1.zzzz;
    
    float4 ixy = permute(permute(ix) + iy);
    float4 ixy0 = permute(ixy + iz0);
    float4 ixy1 = permute(ixy + iz1);
    
    float4 gx0 = ixy0 * (1.0 / 7.0);
    float4 gy0 = fract(floor(gx0) * (1.0 / 7.0)) - 0.5;
    gx0 = fract(gx0);
    float4 gz0 = float4(0.5) - abs(gx0) - abs(gy0);
    float4 sz0 = step(gz0, float4(0.0));
    gx0 -= sz0 * (step(0.0, gx0) - 0.5);
    gy0 -= sz0 * (step(0.0, gy0) - 0.5);
    
    float4 gx1 = ixy1 * (1.0 / 7.0);
    float4 gy1 = fract(floor(gx1) * (1.0 / 7.0)) - 0.5;
    gx1 = fract(gx1);
    float4 gz1 = float4(0.5) - abs(gx1) - abs(gy1);
    float4 sz1 = step(gz1, float4(0.0));
    gx1 -= sz1 * (step(0.0, gx1) - 0.5);
    gy1 -= sz1 * (step(0.0, gy1) - 0.5);
    
    float3 g000 = float3(gx0.x,gy0.x,gz0.x);
    float3 g100 = float3(gx0.y,gy0.y,gz0.y);
    float3 g010 = float3(gx0.z,gy0.z,gz0.z);
    float3 g110 = float3(gx0.w,gy0.w,gz0.w);
    float3 g001 = float3(gx1.x,gy1.x,gz1.x);
    float3 g101 = float3(gx1.y,gy1.y,gz1.y);
    float3 g011 = float3(gx1.z,gy1.z,gz1.z);
    float3 g111 = float3(gx1.w,gy1.w,gz1.w);
    
    float4 norm0 = taylorInvSqrt(float4(dot(g000, g000), dot(g010, g010), dot(g100, g100), dot(g110, g110)));
    g000 *= norm0.x;
    g010 *= norm0.y;
    g100 *= norm0.z;
    g110 *= norm0.w;
    float4 norm1 = taylorInvSqrt(float4(dot(g001, g001), dot(g011, g011), dot(g101, g101), dot(g111, g111)));
    g001 *= norm1.x;
    g011 *= norm1.y;
    g101 *= norm1.z;
    g111 *= norm1.w;
    
    float n000 = dot(g000, Pf0);
    float n100 = dot(g100, float3(Pf1.x, Pf0.yz));
    float n010 = dot(g010, float3(Pf0.x, Pf1.y, Pf0.z));
    float n110 = dot(g110, float3(Pf1.xy, Pf0.z));
    float n001 = dot(g001, float3(Pf0.xy, Pf1.z));
    float n101 = dot(g101, float3(Pf1.x, Pf0.y, Pf1.z));
    float n011 = dot(g011, float3(Pf0.x, Pf1.yz));
    float n111 = dot(g111, Pf1);
    
    float3 fade_xyz = fade(Pf0);
    float4 n_z = mix(float4(n000, n100, n010, n110), float4(n001, n101, n011, n111), fade_xyz.z);
    float2 n_yz = mix(n_z.xy, n_z.zw, fade_xyz.y);
    float n_xyz = mix(n_yz.x, n_yz.y, fade_xyz.x);
    return 2.2 * n_xyz;
}

struct OutEffectsConstants {
    float radius;
    float rotateR;
    float rotateG;
    float rotateB;
    float scatter;
    float width;
    float height;
};

struct VertexOut {
    simd_float4 position [[position]];
    simd_float4 color;
    EffectsConstants effects;
};

vertex VertexOut vertex_main(uint vertexID [[vertex_id]],
                             constant VertexIn *vertices [[buffer(ShaderGradientInputIndexVertices)]],
                             constant ModelConstants &modelConstants [[buffer(ShaderGradientInputIndexModelConstants)]],
                             constant SceneConstants &sceneConstants [[buffer(ShaderGradientInputIndexSceneConstants)]]) {
    VertexIn in = vertices[vertexID];
    VertexOut out;
    float t = sceneConstants.planeModifiers.t;
    float noiseDencity = sceneConstants.planeModifiers.noiseDencity;
    float noiseStrength = sceneConstants.planeModifiers.noiseStrength;
    float outMultiplier = sceneConstants.planeModifiers.outMultiplier;
    float inMultiplier = sceneConstants.planeModifiers.inMultiplier;
    
    float3 normal = float3(sceneConstants.planeModifiers.normalX,
                           sceneConstants.planeModifiers.normalY,
                           sceneConstants.planeModifiers.normalZ);
    
    float4 position = in.position;
    
    float distortion = outMultiplier * cnoise(inMultiplier * position.xyz * noiseDencity + t);
    
    float3 distorted_pos = position.xyz + normal * distortion * noiseStrength;
    
    out.position = sceneConstants.projection * sceneConstants.viewMatrix * modelConstants.modelMatrix * float4(distorted_pos, 1);
    
    if (distorted_pos.z >= 0) {
        float4 diff = sceneConstants.coloring.topColor - in.color;
        float d = saturate(distorted_pos.z / sceneConstants.coloring.topBorder);
        float4 b = diff * d;
        out.color = in.color + b;
    } else {
        float4 diff = sceneConstants.coloring.bottomColor - in.color;
        float d = saturate(distorted_pos.z / sceneConstants.coloring.bottomBorder);
        float4 b = diff * d;
        out.color = in.color + b;
    }
    out.effects = sceneConstants.effects;
    
    return out;
}

// MARK: - Fragment shader

// grid sample positions
struct Cell {
    float2 normal;
    float2 p1;
    float2 p2;
    float2 p3;
    float2 p4;
};

// get pseudo-random number
float rand(float2 seed){
    return fract(sin(dot(seed, float2(12.9898, 78.233))) * 43758.5453);
}

// apply transforms
float distanceToDotRadius( float channel, float2 coord, float2 normal, float2 p, float rad_max) {
    float dist = length(coord - p);
    float rad = pow(channel, 1.125) * rad_max;
    
    return rad - dist;
}

float getDotColour(Cell c, float2 p, float sampleChannel, float angle, float aa, float radius) {
    float dist_c_1 = distanceToDotRadius(sampleChannel, c.p1, c.normal, p, radius);
    float dist_c_2 = distanceToDotRadius(sampleChannel, c.p2, c.normal, p, radius);
    float dist_c_3 = distanceToDotRadius(sampleChannel, c.p3, c.normal, p, radius);
    float dist_c_4 = distanceToDotRadius(sampleChannel, c.p4, c.normal, p, radius);
    
    float res = 0.0;
    res += saturate(dist_c_1 / aa);
    res += saturate(dist_c_2 / aa);
    res += saturate(dist_c_3 / aa);
    res += saturate(dist_c_4 / aa);
    
    return saturate(res);
}

Cell getReferenceCell(float2 p, float grid_angle, float step, float scatter) {
    Cell c;
    
    // calc grid
    float2 n = float2(cos(grid_angle), sin(grid_angle));
    float threshold = step * 0.5;
    float dot_normal = n.x * p.x + n.y * p.y;
    float dot_line = -n.y * p.x + n.x * p.y;
    
    float2 offset = n * dot_normal;
    float offset_normal = fmod(length(offset), step);
    float normal_dir = -1 * sign(dot_normal);
    float normal_scale = (offset_normal < threshold ? -offset_normal : step - offset_normal ) * normal_dir;
    
    float offset_line = fmod(length(float2(p.x - offset.x, p.y - offset.y)), step);
    float line_dir = -1 * sign(dot_line);
    float line_scale = (offset_line < threshold ? -offset_line : step - offset_line ) * line_dir;
    
    // get closest corner
    c.normal = n;
    c.p1.x = p.x - n.x * normal_scale + n.y * line_scale;
    c.p1.y = p.y - n.y * normal_scale - n.x * line_scale;
    
    // scatter
    float off_mag = scatter * threshold * 0.5;
    float off_angle = rand(float2(floor(c.p1.x), floor(c.p1.y))) * (2 * M_PI_F);
    c.p1.x += cos(off_angle) * off_mag;
    c.p1.y += sin(off_angle) * off_mag;
    
    // find corners
    float normal_step = normal_dir * (offset_normal < threshold ? step : -step);
    float line_step = line_dir * (offset_line < threshold ? step : -step);
    c.p2.x = c.p1.x - n.x * normal_step;
    c.p2.y = c.p1.y - n.y * normal_step;
    c.p3.x = c.p1.x + n.y * line_step;
    c.p3.y = c.p1.y - n.x * line_step;
    c.p4.x = c.p1.x - n.x * normal_step + n.y * line_step;
    c.p4.y = c.p1.y - n.y * normal_step - n.x * line_step;
    
    return c;
}

fragment float4 fragment_main(VertexOut in [[stage_in]]) {
    float2 p = in.position.xy;

    // get channel samples
    Cell cell_r = getReferenceCell(p, in.effects.rotateR, in.effects.radius, in.effects.scatter);
    Cell cell_g = getReferenceCell(p, in.effects.rotateG, in.effects.radius, in.effects.scatter);
    Cell cell_b = getReferenceCell(p, in.effects.rotateB, in.effects.radius, in.effects.scatter);

    float4 sampleColor = in.color;
    float aa = ( in.effects.radius < 2.5 ) ? in.effects.radius * 0.5 : 1.25;
    float r = getDotColour(cell_r, p, sampleColor.r, in.effects.rotateR, aa, in.effects.radius);
    float g = getDotColour(cell_g, p, sampleColor.g, in.effects.rotateG, aa, in.effects.radius);
    float b = getDotColour(cell_b, p, sampleColor.b, in.effects.rotateB, aa, in.effects.radius);
    
    // Greyscale
//    r = g = b = (r + b + g) / 3.0;
    
    return float4(r, g, b, 1.0);
}