WaveFunctionCollapser.cs

using System;
using System.Collections.Generic;
using System.Runtime.Intrinsics;
using System.Threading;
using System.Threading.Tasks;
using PhysicsGame.Core.Extensions;

public class WaveFunctionCollapser(IWaveCollapseHandler model, WaveCollapseConfig config)
{
    private IWaveCollapseHandler Model { get; } = model;
    private WaveCollapseConfig Config { get; } = config;

    public async Task Run(CancellationToken cancellationToken)
    {
        while (!cancellationToken.IsCancellationRequested)
        {
            var data = Model.GetCollapeasable(steps: 1);
            var ran = false;
            while (ran |= TryHandleData(data))
            {
            }
            if (!ran)
            {
                await Task.Delay(1000, cancellationToken);
            }
        }
    }

    private const int vL = 0;
    private const int vTL = 1;
    private const int vT = 2;
    private const int vTR = 3;
    private const int vR = 4;
    private const int vBR = 5;
    private const int vB = 6;
    private const int vBL = 7;

    private readonly HashSet<Vector256<uint>> _selfCollapseMemory = [];
    private readonly Dictionary<Vector256<uint>, Vector256<uint>> _trimMemory = [];
    public bool TryHandleData(IWaveCollapseData2D data)
    {
        // Get the lowest entropy next cell
        (int x, int y) best = (-1, -1);
        var lowestEntropy = int.MaxValue;
        for (var x = data.Left; x < data.Right; x++)
        {
            for (var y = data.Top; y < data.Bottom; y++)
            {
                var entropy = data[x,y].HammingWeight();
                // Skip already collapsed cells
                if (entropy == 8 || entropy >= lowestEntropy)
                {
                    continue;
                }
                lowestEntropy = entropy;
                best = (x, y);
            }
        }
        if (lowestEntropy == 1)
        {
            return false;
        }

        // Start collapsing
        Queue<(int x, int y)> collapseQueue = [];
        collapseQueue.Enqueue(best);
        var first = true;
        while (collapseQueue.TryDequeue(out var next))
        {
            var (x, y) = next;
            // Collapsing out of bounds
            if (x < data.Left || x > data.Right || y < data.Top || y > data.Bottom)
            {
                continue;
            }

            // If this was the first cell, or propogating caused a change
            // Then we want to re-queue all adjacent cells for checking
            if (first || Propogate(data, x, y))
            {
                collapseQueue.Enqueue((x - 1, y)); // Left
                collapseQueue.Enqueue((x + 1, y)); // Right
                collapseQueue.Enqueue((x, y - 1)); // Top
                collapseQueue.Enqueue((x, y + 1)); // Bottom
            }
            // If its the first cell of the wave, roll a collapse on it
            if (first)
            {
                Collapse(data, x, y);
                first = false;
            }
        }

        return true;
    }

    // Collapses a cell's superpositional state to one of its weighted
    // Rolled options
    private void Collapse(IWaveCollapseData2D data, int x, int y)
    {
        // Get Weights
        var weights = GetWeights(data[x, y]);
        if (weights.Count == 1)
        {
            Model.Collapse(x, y, weights[0].Option);
        }
        else
        {
            var sum = 0;
            foreach (var (_, Weight) in weights)
            {
                sum += Weight;
            }
            var roll = Random.Shared.Next(1, sum + 1);
            for (var n = 0; n < weights.Count; n++)
            {
                if (roll < weights[n].Weight)
                {
                    Model.Collapse(x, y, weights[n].Option);
                    break;
                }
                roll -= weights[n].Weight;
            }
        }

    }

    private bool Propogate(IWaveCollapseData2D data, int x, int y)
    {
        // Neighbouring cells
        var left = x == data.Left ? Config.OuterBounds : data[x-1, y];
        var right = x == data.Right ? Config.OuterBounds : data[x+1, y];
        var top = y == data.Top ? Config.OuterBounds : data[x, y-1];
        var bottom = y == data.Bottom ? Config.OuterBounds : data[x, y+1];

        var newValue = Vector256.Create(
                left[vR], // Left
                left[vTR] & top[vBL], // Top Left
                top[vB], // Top
                top[vBR] & right[vTL], // Top Right
                right[vL], // Right
                right[vBL] & bottom[vTR], // Bottom Right
                bottom[vT], // Bottom
                bottom[vTL] & left[vBR] // Bottom Left
            );

        // Didnt change
        if (newValue == data[x, y])
        {
            return false;
        }
        // Check our caches
        if (_selfCollapseMemory.Contains(newValue))
        {
            Model.Collapse(x, y, newValue);
        }
        else if (_trimMemory.TryGetValue(newValue, out var trimmedValue))
        {
            Model.Collapse(x, y, trimmedValue);
        }
        // Wasnt cached, cache it
        else
        {
            trimmedValue = TrimData(newValue);
            if (newValue == trimmedValue)
            {
                _ = _selfCollapseMemory.Add(trimmedValue);
            }
            else
            {
                _trimMemory[newValue] = trimmedValue;
            }
            Model.Collapse(x, y, trimmedValue);
        }

        return true;
    }


    private readonly Dictionary<Vector256<uint>, List<(Vector256<uint> Option, int Weight)>> _weightMemory = [];
    // Fetches and caches if neeeded the weighted list for a given superposition
    private List<(Vector256<uint> Option, int Weight)> GetWeights(Vector256<uint> value)
    {
        if (!_weightMemory.TryGetValue(value, out var weights))
        {
            weights = [];
            foreach (var option in GetOptions(value))
            {
                weights.Add((option, Model.Lookup[option]));
            }
            _weightMemory[value] = weights;
        }
        return weights;
    }

    // Fetches the or'd together list of all real possible options for a given
    // super position
    private Vector256<uint> TrimData(Vector256<uint> from)
    {
        var to = Vector256<uint>.Zero;
        // If the option is valid for our from data, add it to our output mapping
        foreach (var option in GetOptions(from))
        {
            to |= option;
        }

        return to;
    }

    // Gets all valid options on the model's lookup chart for a given super position
    // That it is capable of collapsing into validly
    private IEnumerable<Vector256<uint>> GetOptions(Vector256<uint> value)
    {
        foreach (var option in Model.Lookup.Keys)
        {
            if ((value & option) == option)
            {
                yield return option;
            }
        }
    }
}