DigTaskStartSystem.cs

namespace Project {
	using Unity.Burst;
	using Unity.Collections;
	using Unity.Entities;
	using Unity.Jobs;
	using Unity.Mathematics;
	using static Unity.Mathematics.math;

	[UpdateInGroup(typeof(SimulationSystemGroup))]
	public class DigTaskStartSystem : JobComponentSystem {
		public JobHandle dependencies;

		private EndSimulationEntityCommandBufferSystem system;

		protected override void OnCreate() =>
			system = World.GetExistingSystem<EndSimulationEntityCommandBufferSystem>();

		protected override JobHandle OnUpdate(JobHandle dependencies) {
			var agents = GetEntityQuery(new EntityQueryDesc {
				All = new ComponentType[] { typeof(Position) },
				None = new ComponentType[] { typeof(DigTask), typeof(DigInstruction) },
			}).ToEntityArray(Allocator.TempJob);
			var tasks = GetEntityQuery(typeof(Position), typeof(DigInstruction))
				.ToEntityArray(Allocator.TempJob);
			var positions = GetComponentDataFromEntity<Position>(isReadOnly: true);
			var weights = new NativeArray<float>(tasks.Length * agents.Length, Allocator.TempJob, NativeArrayOptions.UninitializedMemory);
			dependencies = new FloatMaxJob {
				floats = weights,
			}.Schedule(weights.Length, 1, dependencies);
			dependencies = new PathJob {
				width = Game.Width,
				height = Game.Height,
				cells = Game.Cells,
				tasks = tasks,
				agents = agents,
				positions = positions,
				weights = weights,
			}.Schedule(tasks.Length, 1, dependencies);
			dependencies = new ReduceColumnJob {
				width = tasks.Length,
				height = agents.Length,
				weights = weights,
			}.Schedule(tasks.Length, 1, dependencies);
			dependencies = new ReduceRowJob {
				width = tasks.Length,
				weights = weights,
			}.Schedule(agents.Length, 1, dependencies);
			var E = new NativeArray<bool>(tasks.Length * agents.Length, Allocator.TempJob, NativeArrayOptions.UninitializedMemory);
			dependencies = new MarkElligibleJob {
				weights = weights,
				E = E,
			}.Schedule(E.Length, 1, dependencies);
			var M = new NativeArray<bool>(tasks.Length * agents.Length, Allocator.TempJob);
			dependencies = new MaximalMatchingJob {
				width = tasks.Length,
				height = agents.Length,
				E = E,
				M = M,
			}.Schedule(dependencies);
			var buffer = system.CreateCommandBuffer().ToConcurrent();
			var taskCount = tasks.Length;
			var agentCount = agents.Length;
			dependencies = Entities
				.WithName("AssignJob")
				.WithReadOnly(agents)
				.WithReadOnly(M)
				.WithDeallocateOnJobCompletion(agents)
				.WithDeallocateOnJobCompletion(M)
				.WithDeallocateOnJobCompletion(tasks)
				.WithAll<DigInstruction>()
				.ForEach((Entity task, int entityInQueryIndex, in Position position) => {
					tasks[0] = tasks[0];
					for (var agent = 0; agent < agentCount; ++agent) {
						if (M[entityInQueryIndex + agent * taskCount]) {
							buffer.DestroyEntity(entityInQueryIndex, task);
							buffer.AddComponent(entityInQueryIndex, agents[agent], new DigTask {
								progress = 0f,
							});
							buffer.AddComponent(entityInQueryIndex, agents[agent], new Path {
								position = (int2)position.value,
								range = 1,
							});
							break;
						}
					}
				})
				.Schedule(dependencies);
			system.AddJobHandleForProducer(dependencies);
			return this.dependencies = dependencies;
		}

		[BurstCompile]
		private struct FloatMaxJob : IJobParallelFor {
			[WriteOnly] public NativeArray<float> floats;

			public void Execute(int index) {
				floats[index] = float.MaxValue;
			}
		}

		[BurstCompile]
		private struct PathJob : IJobParallelFor {
			public int width;
			public int height;
			[ReadOnly] public ComponentDataFromEntity<Position> positions;
			[ReadOnly] public NativeArray<Cell> cells;
			[ReadOnly] public NativeArray<Entity> tasks;
			[ReadOnly] public NativeArray<Entity> agents;
			[WriteOnly, NativeDisableParallelForRestriction] public NativeArray<float> weights;

			//   x
			//  012
			// y3_5z
			//  678
			//   w
			private bool3x3 Prune(int2 p, bool4 n) {
				if (p.Equals(int2(-1, -1))) { // Moving down and right
					return new bool3x3(
						  true,   true,    n.z,
						  true,  false,  false,
						   n.w,  false,  false
					);
				}
				if (p.Equals(int2(-1, 0))) { // Moving right
					return new bool3x3(
						   n.x,  false,  false,
						  true,  false,  false,
						   n.w,  false,  false
					);
				}
				if (p.Equals(int2(-1, +1))) { // Moving up and right
					return new bool3x3(
						   n.x,  false,  false,
						  true,  false,  false,
						  true,   true,    n.z
					);
				}
				if (p.Equals(int2(0, -1))) { // Moving down
					return new bool3x3(
						   n.y,   true,    n.z,
						 false,  false,  false,
						 false,  false,  false
					);
				}
				if (p.Equals(int2(0, +1))) { // Moving up
					return new bool3x3(
						  false,  false,  false,
						  false,  false,  false,
						    n.y,   true,    n.z
					);
				}
				if (p.Equals(int2(+1, -1))) { // Moving down and left
					return new bool3x3(
						   n.y,   true,   true,
						 false,  false,   true,
						 false,  false,    n.w
					);
				}
				if (p.Equals(int2(+1, 0))) { // Moving left
					return new bool3x3(
						 false,  false,    n.x,
						 false,  false,   true,
						 false,  false,    n.w
					);
				}
				if (p.Equals(int2(+1, +1))) { // Moving up and left
					return new bool3x3(
						 false,  false,    n.x,
						 false,  false,   true,
						   n.y,   true,   true
					);
				}
				return new bool3x3(true, true, true, true, false, true, true, true, true);
			}

			// TODO: Implement properly
			private bool Forced(int2 n) {
				for (var y = -1; y <= 1; ++y) {
					for (var x = -1; x <= 1; ++x) {
						if (x == 0 && y != 0 || x != 0 && y == 0) {
							if ((cells[n.x + x + (n.y + y) * width].flags & Flags.Collision) != Flags.None) {
								return true;
							}
						}
					}
				}
				return false;
			}

			private int2 Jump(int2 x, int2 d, int2 g) {
				var n = x + d;
				if (any(n < 0) || any(n > int2(width, height)) || (cells[n.x + n.y * width].flags & Flags.Collision) != Flags.None) { 
					return int2(-1, -1);
				}
				if (n.Equals(g)) {
					return n;
				}
				if (Forced(n)) {
					return n;
				}
				if (all(abs(d))) {
					if (!Jump(n, int2(d.x, 0), g).Equals(int2(-1, -1))) {
						return n;
					}
					if (!Jump(n, int2(0, d.y), g).Equals(int2(-1, -1))) {
						return n;
					}
				}
				return Jump(n, d, g);
			}

			public void Execute(int task) {
				var open = new NativeList<int2>(Allocator.Temp);
				var g = new NativeHashMap<int2, float>(1, Allocator.Temp);
				var f = new NativeHashMap<int2, float>(1, Allocator.Temp);
				var p = new NativeHashMap<int2, int2>(1, Allocator.Temp);
				var start = (int2)positions[tasks[task]].value;
				for (var agent = 0; agent < agents.Length; ++agent) {
					var agentPosition = positions[agents[agent]].value;
					var goal = (int2)agentPosition;
					open.Clear();
					g.Clear();
					f.Clear();
					p.Clear();
					open.Add(start);
					g[start] = 0f;
					f[start] = abs(start.x - goal.x) + abs(start.y - goal.y);
					p[start] = int2(0, 0);

					while (open.Length > 0) {
						var index = open.Length - 1;
						var current = open[index];
						var minimum = f[current];
						for (var i = index - 1; i >= 0; --i) {
							var candidate = open[i];
							var estimate = f[candidate];
							if (estimate < minimum) {
								minimum = estimate;
								index = i;
							}
						}
						current = open[index];
						var currentID = current.x + current.y * width;
						open.RemoveAtSwapBack(index);

						// Done?
						if (current.Equals(goal)) {
							weights[task + agent * tasks.Length] = g[current];
							break;
						}

						var neighbors = Prune(p[current], new bool4(
							(cells[currentID + 0 - width].flags & Flags.Collision) != Flags.None,
							(cells[currentID - 1 +     0].flags & Flags.Collision) != Flags.None,
							(cells[currentID + 1 +     0].flags & Flags.Collision) != Flags.None,
							(cells[currentID + 0 + width].flags & Flags.Collision) != Flags.None
						));
						
						var currentG = g[current];
						for (var y = -1; y <= 1; ++y) {
							for (var x = -1; x <= 1; ++x) {
								if (!neighbors[x + 1][y + 1]) { continue; }
								var direction = int2(x, y);
								var next = Jump(current, direction, goal);
								if (all(next > -1)) {
									var d = abs(current.x - next.x) + abs(current.y - next.y);
									var tentativeG = currentG + d;
									if (!g.ContainsKey(next) || tentativeG < g[next]) {
										var h = abs(goal.x - next.x) + abs(goal.y - next.y);
										g[next] = tentativeG;
										f[next] = tentativeG + h;
										p[next] = direction;
										if (!open.Contains(next)) {
											open.Add(next);
										}
									}
								}
							}
						}
					}
				}
				g.Dispose();
				f.Dispose();
				p.Dispose();
				open.Dispose();
			}
		}

		[BurstCompile]
		private struct ReduceColumnJob : IJobParallelFor {
			public int width;
			public int height;
			[NativeDisableParallelForRestriction] public NativeArray<float> weights;

			public void Execute(int i) {
				var min = float.MaxValue;
				for (var j = 0; j < height; ++j) {
					var weight = weights[i + j * width];
					if (weight < min) {
						min = weight;
					}
				}
				if (min < float.MaxValue) {
					for (var j = 0; j < height; ++j) {
						weights[i + j * width] -= min;
					}
				}
			}
		}

		[BurstCompile]
		private struct ReduceRowJob : IJobParallelFor {
			public int width;
			[NativeDisableParallelForRestriction] public NativeArray<float> weights;

			public void Execute(int j) {
				var min = float.MaxValue;
				for (var i = 0; i < width; ++i) {
					var weight = weights[i + j * width];
					if (weight < min) {
						min = weight;
					}
				}
				if (min < float.MaxValue) {
					for (var i = 0; i < width; ++i) {
						weights[i + j * width] -= min;
					}
				}
			}
		}

		[BurstCompile]
		private struct MarkElligibleJob : IJobParallelFor {
			[ReadOnly, DeallocateOnJobCompletion] public NativeArray<float> weights;
			[WriteOnly, NativeDisableParallelForRestriction] public NativeArray<bool> E;

			public void Execute(int index) {
				E[index] = weights[index] == 0f;
			}
		}

		[BurstCompile]
		private struct MaximalMatchingJob : IJob {
			public int width;
			public int height;
			[ReadOnly, DeallocateOnJobCompletion] public NativeArray<bool> E;
			[WriteOnly] public NativeArray<bool> M;

			public void Execute() {
				var C = new NativeArray<bool>(width, Allocator.Temp);
				var R = new NativeArray<bool>(height, Allocator.Temp);
				for (var y = 0; y < height; ++y) {
					for (var x = 0; x < width; ++x) {
						var index = x + y * width;
						if (E[index] && !C[x] && !R[y]) {
							M[index] = true;
							C[x] = true;
							R[y] = true;
							break;
						}
					}
				}
				C.Dispose();
				R.Dispose();
			}
		}
	}
}