1// Copyright 2020 The Tint Authors.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15// vertex shader
16
17[[stage(vertex)]]
18fn vert_main([[location(0)]] a_particlePos : vec2<f32>,
19             [[location(1)]] a_particleVel : vec2<f32>,
20             [[location(2)]] a_pos : vec2<f32>)
21          -> [[builtin(position)]] vec4<f32> {
22  var angle : f32 = -atan2(a_particleVel.x, a_particleVel.y);
23  var pos : vec2<f32> = vec2<f32>(
24      (a_pos.x * cos(angle)) - (a_pos.y * sin(angle)),
25      (a_pos.x * sin(angle)) + (a_pos.y * cos(angle)));
26  return vec4<f32>(pos + a_particlePos, 0.0, 1.0);
27}
28
29// fragment shader
30
31[[stage(fragment)]]
32fn frag_main() -> [[location(0)]] vec4<f32> {
33  return vec4<f32>(1.0, 1.0, 1.0, 1.0);
34}
35
36// compute shader
37struct Particle {
38  pos : vec2<f32>;
39  vel : vec2<f32>;
40};
41
42[[block]] struct SimParams {
43  deltaT : f32;
44  rule1Distance : f32;
45  rule2Distance : f32;
46  rule3Distance : f32;
47  rule1Scale : f32;
48  rule2Scale : f32;
49  rule3Scale : f32;
50};
51
52[[block]] struct Particles {
53  particles : array<Particle, 5>;
54};
55
56[[binding(0), group(0)]] var<uniform> params : SimParams;
57[[binding(1), group(0)]] var<storage, read_write> particlesA : Particles;
58[[binding(2), group(0)]] var<storage, read_write> particlesB : Particles;
59
60// https://github.com/austinEng/Project6-Vulkan-Flocking/blob/master/data/shaders/computeparticles/particle.comp
61[[stage(compute), workgroup_size(1)]]
62fn comp_main(
63  [[builtin(global_invocation_id)]] gl_GlobalInvocationID : vec3<u32>) {
64  var index : u32 = gl_GlobalInvocationID.x;
65  if (index >= 5u) {
66    return;
67  }
68
69  var vPos : vec2<f32> = particlesA.particles[index].pos;
70  var vVel : vec2<f32> = particlesA.particles[index].vel;
71
72  var cMass : vec2<f32> = vec2<f32>(0.0, 0.0);
73  var cVel : vec2<f32> = vec2<f32>(0.0, 0.0);
74  var colVel : vec2<f32> = vec2<f32>(0.0, 0.0);
75  var cMassCount : i32 = 0;
76  var cVelCount : i32 = 0;
77
78  var pos : vec2<f32>;
79  var vel : vec2<f32>;
80  for(var i : u32 = 0u; i < 5u; i = i + 1u) {
81    if (i == index) {
82      continue;
83    }
84
85    pos = particlesA.particles[i].pos.xy;
86    vel = particlesA.particles[i].vel.xy;
87
88    if (distance(pos, vPos) < params.rule1Distance) {
89      cMass = cMass + pos;
90      cMassCount = cMassCount + 1;
91    }
92    if (distance(pos, vPos) < params.rule2Distance) {
93      colVel = colVel - (pos - vPos);
94    }
95    if (distance(pos, vPos) < params.rule3Distance) {
96      cVel = cVel + vel;
97      cVelCount = cVelCount + 1;
98    }
99  }
100  if (cMassCount > 0) {
101    cMass =
102      (cMass / vec2<f32>(f32(cMassCount), f32(cMassCount))) - vPos;
103  }
104  if (cVelCount > 0) {
105    cVel = cVel / vec2<f32>(f32(cVelCount), f32(cVelCount));
106  }
107
108  vVel = vVel + (cMass * params.rule1Scale) + (colVel * params.rule2Scale) +
109      (cVel * params.rule3Scale);
110
111  // clamp velocity for a more pleasing simulation
112  vVel = normalize(vVel) * clamp(length(vVel), 0.0, 0.1);
113
114  // kinematic update
115  vPos = vPos + (vVel * params.deltaT);
116
117  // Wrap around boundary
118  if (vPos.x < -1.0) {
119    vPos.x = 1.0;
120  }
121  if (vPos.x > 1.0) {
122    vPos.x = -1.0;
123  }
124  if (vPos.y < -1.0) {
125    vPos.y = 1.0;
126  }
127  if (vPos.y > 1.0) {
128    vPos.y = -1.0;
129  }
130
131  // Write back
132  particlesB.particles[index].pos = vPos;
133  particlesB.particles[index].vel = vVel;
134}
135