bluflame/intromat/Intromat/Nodes/Textures/Shape2DGrayscaleNode.cs

using System;
using System.Diagnostics.CodeAnalysis;
using DynamicData;
using Intromat.Model.Compiler;
using Intromat.PersistentModel;
using Intromat.ViewModels;
using Intromat.ViewModels.Editors;
using Intromat.Views;
using NodeNetwork.Toolkit.ValueNode;
using NodeNetwork.ViewModels;
using ReactiveUI;
using SharpDX.D3DCompiler;
using SharpDX.Direct3D11;
using Splat;
using Xceed.Wpf.Toolkit.PropertyGrid.Attributes;
using Buffer = SharpDX.Direct3D11.Buffer;

namespace Intromat.Nodes.Textures
{
    [CategoryOrder("Shape 2D", 1)]
    public class Shape2DGrayscaleNode : DxTextureNodeBase
    {
        private Buffer? _constantBuffer;
        private ComputeShader? _cs;
        private const string _computeShader = @"cbuffer cb
        {
            int t;
            int p;
            float s;
            float sx;
            float sy;
            float a;
            int r;
        }

        float Square(float2 uv)
        {
            return step(uv.x, sx) * step(-sx, uv.x) * step(uv.y, sy) * step(-sy, uv.y);
        }

        RWTexture2D<float> o;
        [numthreads(16,16,1)]
        void main(in uint3 i : SV_DispatchThreadID)
        { 
            float sa, ca;
            sincos(a, sa, ca);
            float2 dim;
            o.GetDimensions(dim.x, dim.y);
            float2 uv = i.xy / dim;
            if (r > 0)
            {
                float f = sqrt(2)/2;
                uv = uv * 2 - 1;
                uv = mul(uv, float2x2(f, -f, f, f));
                uv = uv * .5 + .5;
            }
            uv *= t;
            uv = frac(uv);
            uv = uv * 2 - 1;
            uv /= s;
            uv = mul(uv, float2x2(ca, -sa, sa, ca));
            float z = 0;
            switch (p)
            {
                case 0: z = Square(uv); break;
            }
            o[i.xy] = z;
        }";

        static Shape2DGrayscaleNode()
        {
            Locator.CurrentMutable.Register(() => new CodeGenNodeView(), typeof(IViewFor<Shape2DGrayscaleNode>));
        }

        public Shape2DGrayscaleNode()
        {
            Name = "Shape 2D";
            var group = new EndpointGroup("Shape 2D");
            Inputs.Add(Tiling = new CodeGenInputViewModel<ITypedExpression<int>>(EPortType.Integer)
            {
                Name = "Tiling",
                Group = group,
                Editor = TilingEditor
            });
            Inputs.Add(Pattern = new CodeGenInputViewModel<int>(EPortType.None)
            {
                Name = "Pattern",
                Group = group,
                Editor = PatternEditor
            });
            Inputs.Add(Scale = new CodeGenInputViewModel<ITypedExpression<float>>(EPortType.Float)
            {
                Name = "Scale",
                Group = group,
                Editor = ScaleEditor
            });
            Inputs.Add(SizeX = new CodeGenInputViewModel<ITypedExpression<float>>(EPortType.Float)
            {
                Name = "Size X",
                Group = group,
                Editor = SizeXEditor
            });
            Inputs.Add(SizeY = new CodeGenInputViewModel<ITypedExpression<float>>(EPortType.Float)
            {
                Name = "Size Y",
                Group = group,
                Editor = SizeYEditor
            });
            Inputs.Add(Angle = new CodeGenInputViewModel<ITypedExpression<float>>(EPortType.Float)
            {
                Name = "Angle",
                Group = group,
                Editor = AngleEditor
            });
            Inputs.Add(Rotate45 = new CodeGenInputViewModel<ITypedExpression<bool>>(EPortType.Float)
            {
                Name = "Rotate 45°",
                Group = group,
                Editor = Rotate45Editor
            });
        }

        public override NodeModelBase CreateModel()
        {
            return new Shape2DGrayscaleModel();
        }

        public override void SaveModel(NodeModelBase model)
        {
            base.SaveModel(model);
            var shape2D = (Shape2DGrayscaleModel)model;
            shape2D.Tiling = TilingEditor.CreateModel();
            shape2D.Pattern = (Shape2DGrayscaleModel.EPattern)PatternEditor.Value;
            shape2D.Scale = ScaleEditor.CreateModel();
            shape2D.SizeX = SizeXEditor.CreateModel();
            shape2D.SizeY = SizeYEditor.CreateModel();
            shape2D.Angle = AngleEditor.CreateModel();
            shape2D.Rotate45 = Rotate45Editor.CreateModel();
        }

        public override void LoadModel(NodeModelBase model)
        {
            base.LoadModel(model);
            var shape2D = (Shape2DGrayscaleModel)model;
            TilingEditor.LoadModel(shape2D.Tiling);
            PatternEditor.Value = (int)shape2D.Pattern;
            ScaleEditor.LoadModel(shape2D.Scale);
            SizeXEditor.LoadModel(shape2D.SizeX);
            SizeYEditor.LoadModel(shape2D.SizeY);
            AngleEditor.LoadModel(shape2D.Angle);
            Rotate45Editor.LoadModel(shape2D.Rotate45);
        }

        protected override unsafe void UpdateFrame(Device device, DeviceContext context)
        {
            base.UpdateFrame(device, context);

            context.MapSubresource(_constantBuffer, 0, MapMode.WriteDiscard, MapFlags.None, out var stream);
            var floats = new Span<float>((void*)stream.DataPointer, 8);
            var ints = new Span<int>((void*)stream.DataPointer, 8);
            ints[0] = Tiling.Value.Evaluate();
            ints[1] = Pattern.Value;
            floats[2] = Scale.Value.Evaluate();
            floats[3] = SizeX.Value.Evaluate();
            floats[4] = SizeY.Value.Evaluate();
            floats[5] = Angle.Value.Evaluate();
            ints[6] = Rotate45.Value.Evaluate() ? 1 : 0;
            context.UnmapSubresource(_constantBuffer, 0);

            context.ComputeShader.SetConstantBuffer(0, _constantBuffer);
            context.ComputeShader.SetShader(_cs, null, 0);
            context.ComputeShader.SetUnorderedAccessView(0, _output!.UnorderedAccessView);
            context.Dispatch(_texDesc.Width / 16, _texDesc.Height / 16, 1);
            context.ComputeShader.SetUnorderedAccessView(0, null);
        }

        [MemberNotNull(nameof(_constantBuffer))]
        [MemberNotNull(nameof(_cs))]
        protected override void CreateDeviceResources(Device device)
        {
            base.CreateDeviceResources(device);

            var bufferDesc = new BufferDescription(32, ResourceUsage.Dynamic, BindFlags.ConstantBuffer, CpuAccessFlags.Write, ResourceOptionFlags.None, 0);
            _constantBuffer = new Buffer(device, bufferDesc);
            var flags = ShaderFlags.None;
#if DEBUG
            flags |= ShaderFlags.Debug;
#endif
            _cs = new ComputeShader(device, ShaderBytecode.Compile(_computeShader, "main", "cs_5_0", flags).Bytecode.Data);
        }

        public IntegerExpressionEditorViewModel TilingEditor { get; } = new() { CustomValue = 1, MinValue = 1, MaxValue = 16 };
        public ValueNodeInputViewModel<ITypedExpression<int>> Tiling { get; }

        public EnumEditorViewModel PatternEditor { get; } = new(typeof(Shape2DGrayscaleModel.EPattern));
        public ValueNodeInputViewModel<int> Pattern { get; }

        public FloatExpressionEditorViewModel ScaleEditor { get; } = new() { CustomValue = 1, MaxValue = 1 };
        public ValueNodeInputViewModel<ITypedExpression<float>> Scale { get; }

        public FloatExpressionEditorViewModel SizeXEditor { get; } = new() { CustomValue = .5f, MaxValue = 1 };
        public ValueNodeInputViewModel<ITypedExpression<float>> SizeX { get; }

        public FloatExpressionEditorViewModel SizeYEditor { get; } = new() { CustomValue = .5f, MaxValue = 1 };
        public ValueNodeInputViewModel<ITypedExpression<float>> SizeY { get; }

        public FloatExpressionEditorViewModel AngleEditor { get; } = new() { MaxValue = (float)(2 * Math.PI) };
        public ValueNodeInputViewModel<ITypedExpression<float>> Angle { get; }

        public BooleanExpressionEditorViewModel Rotate45Editor { get; } = new();
        public ValueNodeInputViewModel<ITypedExpression<bool>> Rotate45 { get; }
    }
}