fleeting-garden/src/pipelines/agents/agent.wgsl

struct Settings {
  moveRate: f32,
  turnRate: f32,
  sensorAngleSin: f32,
  sensorAngleCos: f32,
  sensorOffset: f32,
  turnWhenLost: f32,
  individualTrailWeight: f32,
  agentCount: f32,
  introProgress: f32,
  color1ToColor1: f32,
  color1ToColor2: f32,
  color1ToColor3: f32,
  color2ToColor1: f32,
  color2ToColor2: f32,
  color2ToColor3: f32,
  color3ToColor1: f32,
  color3ToColor2: f32,
  color3ToColor3: f32,
  sourceAttractionWeight: f32,
  sourceSlowMoveRate: f32,
  sourceTrailWeightMultiplier: f32,
  forwardRotationScale: f32,
  introNearDistanceInner: f32,
  introNearDistanceMin: f32,
  introNearSensorOffsetMultiplier: f32,
  introTargetAngleBlend: f32,
  introProgressCutoff: f32,
  introTurnRateMultiplier: f32,
  introRandomTurnMultiplier: f32,
  introFarMoveMultiplier: f32,
  introNearMoveMultiplier: f32,
  introStepStopDistance: f32,
  randomTimeScale: f32,
};

@group(1) @binding(0) var<uniform> settings: Settings;
@group(1) @binding(2) var trailMapIn: texture_2d<f32>;
@group(1) @binding(3) var trailMapOut: texture_storage_2d<rgba16float, write>;
@group(1) @binding(4) var sourceMap: texture_2d<f32>;

@compute @workgroup_size(64)
fn main(
  @builtin(global_invocation_id) global_id: vec3<u32>,
  @builtin(num_workgroups) num_workgroups: vec3<u32>
) {
  let id = get_id(global_id, num_workgroups);

  if id >= u32(settings.agentCount) {
    return;
  }

  let colorIndex = agents[id].colorIndex;
  if colorIndex < 0.0 || colorIndex >= 2.5 {
    return;
  }

  var position = agents[id].position;
  var angle = agents[id].angle;
  var targetPosition = vec2<f32>(-1.0, -1.0);
  var hasIntroTarget = false;
  if settings.introProgress < settings.introProgressCutoff {
    targetPosition = agents[id].targetPosition;
    hasIntroTarget = targetPosition.x >= 0.0 && targetPosition.y >= 0.0;
    if hasIntroTarget && settings.introProgress < agents[id].introDelay {
      return;
    }
  }

  let randomSeed = random_seed(id, state.time);
  let randomTurn = random_float(randomSeed);
  let direction = vec2(cos(angle), sin(angle));

  let forwardSensor = sensor_position(position, direction, settings.sensorOffset);
  let leftSensor = sensor_position(
    position,
    rotate_direction(direction, settings.sensorAngleSin, settings.sensorAngleCos),
    settings.sensorOffset
  );
  let rightSensor = sensor_position(
    position,
    rotate_direction(direction, -settings.sensorAngleSin, settings.sensorAngleCos),
    settings.sensorOffset
  );

  let trailForward = textureLoad(trailMapIn, forwardSensor, 0);
  let trailLeft = textureLoad(trailMapIn, leftSensor, 0);
  let trailRight = textureLoad(trailMapIn, rightSensor, 0);
  let sourceForwardSample = textureLoad(sourceMap, forwardSensor, 0);
  let sourceLeftSample = textureLoad(sourceMap, leftSensor, 0);
  let sourceRightSample = textureLoad(sourceMap, rightSensor, 0);

  let channelMask = get_channel_mask(colorIndex);
  let reactionMask = get_reaction_mask(colorIndex);

  let trailForwardWeight = dot(trailForward.rgb, reactionMask);
  let trailLeftWeight = dot(trailLeft.rgb, reactionMask);
  let trailRightWeight = dot(trailRight.rgb, reactionMask);

  let sourceForwardWeight = dot(sourceForwardSample.rgb, reactionMask);
  let sourceLeftWeight = dot(sourceLeftSample.rgb, reactionMask);
  let sourceRightWeight = dot(sourceRightSample.rgb, reactionMask);

  let weightForward =
    trailForwardWeight + sourceForwardWeight * settings.sourceAttractionWeight;
  let weightLeft = trailLeftWeight + sourceLeftWeight * settings.sourceAttractionWeight;
  let weightRight =
    trailRightWeight + sourceRightWeight * settings.sourceAttractionWeight;

  var rotation = (randomTurn - 0.5) * settings.turnWhenLost;
  if weightForward >= weightLeft && weightForward >= weightRight {
    rotation = rotation * settings.forwardRotationScale;
  } else {
    rotation += sign(weightLeft - weightRight) * settings.turnRate;
  }

  let sourceAtAgent = textureLoad(sourceMap, vec2<i32>(position), 0);
  let positiveReactionMask = max(reactionMask, vec3<f32>(0.0));
  let sourceAtAgentStrength = clamp(dot(sourceAtAgent.rgb, positiveReactionMask), 0.0, 1.0);
  var moveRate = settings.moveRate * mix(1.0, settings.sourceSlowMoveRate, sourceAtAgentStrength);
  var introTargetOffset = vec2<f32>(0.0, 0.0);
  var introTargetDistance = 0.0;

  if hasIntroTarget {
    introTargetOffset = targetPosition - position;
    introTargetDistance = length(introTargetOffset);
    let targetAngle = atan2(introTargetOffset.y, introTargetOffset.x);
    let nearTitle = 1.0 - smoothstep(
      settings.introNearDistanceInner,
      max(
        settings.introNearDistanceMin,
        settings.sensorOffset * settings.introNearSensorOffsetMultiplier
      ),
      introTargetDistance
    );
    let desiredAngle = mix(
      targetAngle,
      agents[id].targetAngle,
      nearTitle * settings.introTargetAngleBlend
    );
    let introTurn = angle_delta(angle, desiredAngle);

    rotation = clamp(
        introTurn,
        -settings.turnRate * settings.introTurnRateMultiplier,
        settings.turnRate * settings.introTurnRateMultiplier
      )
      + (random_float(randomSeed + 1013904223u) - 0.5) *
        settings.turnWhenLost *
        settings.introRandomTurnMultiplier;
    moveRate = min(
      settings.moveRate *
        mix(settings.introFarMoveMultiplier, settings.introNearMoveMultiplier, nearTitle),
      introTargetDistance
    );
  }

  var step = direction * moveRate;
  if hasIntroTarget {
    step = vec2<f32>(0.0, 0.0);
    if introTargetDistance > settings.introStepStopDistance {
      step = introTargetOffset / introTargetDistance * moveRate;
    }
  }

  let maxPosition = state.size - vec2<f32>(1.0, 1.0);
  let nextPosition = clamp(position + step, vec2<f32>(0, 0), maxPosition);
  if nextPosition.x == 0 || nextPosition.x == maxPosition.x || nextPosition.y == 0 || nextPosition.y == maxPosition.y {
    rotation = 3.14159265359 + random_float(randomSeed + 22695477u) - 0.5;
  }

  let sourceBelow = textureLoad(sourceMap, vec2<i32>(nextPosition), 0);
  let sourceBelowStrength = clamp(dot(sourceBelow.rgb, positiveReactionMask), 0.0, 1.0);
  let trailWeight =
    settings.individualTrailWeight *
    (1.0 + sourceBelowStrength * settings.sourceTrailWeightMultiplier);
  var trailBelow = textureLoad(trailMapIn, vec2<i32>(nextPosition), 0);
  trailBelow = vec4<f32>(
    trailBelow.rgb + channelMask * trailWeight,
    max(trailBelow.a, 0.0)
  );

  textureStore(trailMapOut, vec2<i32>(nextPosition), trailBelow);
  agents[id].angle = angle + rotation;
  agents[id].position = nextPosition;
}

fn sensor_position(agentPosition: vec2<f32>, direction: vec2<f32>, sensorOffset: f32) -> vec2<i32> {
  return vec2<i32>(clamp(
    agentPosition + direction * sensorOffset,
    vec2<f32>(0, 0),
    state.size - vec2<f32>(1, 1)
  ));
}

fn rotate_direction(direction: vec2<f32>, angleSin: f32, angleCos: f32) -> vec2<f32> {
  return vec2<f32>(
    direction.x * angleCos - direction.y * angleSin,
    direction.x * angleSin + direction.y * angleCos
  );
}

fn get_channel_mask(colorIndex: f32) -> vec3<f32> {
  if colorIndex < 0.5 {
    return vec3<f32>(1, 0, 0);
  }
  if colorIndex < 1.5 {
    return vec3<f32>(0, 1, 0);
  }
  if colorIndex < 2.5 {
    return vec3<f32>(0, 0, 1);
  }
  return vec3<f32>(0.0, 0.0, 0.0);
}

fn get_reaction_mask(colorIndex: f32) -> vec3<f32> {
  if colorIndex < 0.5 {
    return vec3<f32>(
      settings.color1ToColor1,
      settings.color1ToColor2,
      settings.color1ToColor3
    );
  }
  if colorIndex < 1.5 {
    return vec3<f32>(
      settings.color2ToColor1,
      settings.color2ToColor2,
      settings.color2ToColor3
    );
  }
  if colorIndex < 2.5 {
    return vec3<f32>(
      settings.color3ToColor1,
      settings.color3ToColor2,
      settings.color3ToColor3
    );
  }
  return vec3<f32>(0.0, 0.0, 0.0);
}

fn angle_delta(sourceAngle: f32, targetAngle: f32) -> f32 {
  return atan2(sin(targetAngle - sourceAngle), cos(targetAngle - sourceAngle));
}

fn random_seed(id: u32, time: f32) -> u32 {
  let timeSeed = u32(time * settings.randomTimeScale);
  return id * 747796405u + timeSeed * 2891336453u;
}

fn random_float(seed: u32) -> f32 {
  return f32(hash_u32(seed) >> 8u) * (1.0 / 16777216.0);
}

fn hash_u32(seed: u32) -> u32 {
  let value = seed * 747796405u + 2891336453u;
  let word = ((value >> ((value >> 28u) + 4u)) ^ value) * 277803737u;
  return (word >> 22u) ^ word;
}