slime-moldies/slime_simulation.gd

extends Node2D

@export var texture_size: Vector2i = Vector2i(512, 512)
@export var decay_factor: float = 0.98

@onready var mesh := $MeshInstance2D

var texture: Texture2DRD
var next_texture: int = 0


func _ready():
	RenderingServer.call_on_render_thread(_initialize_compute_code.bind(texture_size))
	var mat: ShaderMaterial = mesh.material
	texture = mat.get_shader_parameter("trail_tex")


func _process(_delta: float) -> void:
	next_texture = (next_texture + 1) % 2

	if texture:
		texture.texture_rd_rid = texture_rds[next_texture]

	RenderingServer.call_on_render_thread(
		_render_process.bind(next_texture, texture_size, decay_factor)
	)


func _exit_tree() -> void:
	if texture:
		texture.texture_rd_id = RID()

	RenderingServer.call_on_render_thread(_free_compute_resources)


# Everything after this point is designed to run on our rendering thread.

var rd: RenderingDevice

var decay_shader: RID
var decay_pipeline: RID

var agent_shader: RID
var agent_pipeline: RID

# We use 3 textures:
# - One to render into
# - One that contains the last frame rendered
# - One for the frame before that
var texture_rds: Array[RID] = [RID(), RID()]
var texture_sets: Array[RID] = [RID(), RID()]

var agent_uniform_sets: Array[RID] = [RID(), RID()]
var agent_buffer: RID
var agent_set: RID


func _load_shader(rd_: RenderingDevice, path: String) -> RID:
	var shader_file := load(path)
	var shader_spirv: RDShaderSPIRV = shader_file.get_spirv()
	return rd_.shader_create_from_spirv(shader_spirv)


func _create_trail_uniform(texture_rd: RID) -> RDUniform:
	var uniform := RDUniform.new()
	uniform.uniform_type = RenderingDevice.UNIFORM_TYPE_IMAGE
	uniform.binding = 0
	uniform.add_id(texture_rd)
	return uniform


func _initialize_compute_code(init_with_texture_size: Vector2i) -> void:
	# As this becomes part of our normal frame rendering,
	# we use our main rendering device here.
	rd = RenderingServer.get_rendering_device()

	# Create shaders
	decay_shader = _load_shader(rd, "res://shaders/decay_cs.glsl")
	decay_pipeline = rd.compute_pipeline_create(decay_shader)

	agent_shader = _load_shader(rd, "res://shaders/agents_cs.glsl")
	agent_pipeline = rd.compute_pipeline_create(agent_shader)

	# Create agent buffer
	var rng := RandomNumberGenerator.new()
	rng.seed = 80085
	var agent_data := PackedVector2Array()
	for n in 64 * 16:
		var new_position: Vector2i = Vector2(
			rng.randi_range(0, init_with_texture_size.x),
			rng.randi_range(0, init_with_texture_size.y)
		)
		agent_data.push_back(new_position)

	var agent_size_in_bytes := 8
	agent_buffer = rd.storage_buffer_create(
		agent_data.size() * agent_size_in_bytes, agent_data.to_byte_array()
	)

	# Create trail texture
	var tf: RDTextureFormat = RDTextureFormat.new()
	tf.format = RenderingDevice.DATA_FORMAT_R32_SFLOAT
	tf.texture_type = RenderingDevice.TEXTURE_TYPE_2D
	tf.width = init_with_texture_size.x
	tf.height = init_with_texture_size.y
	tf.depth = 1
	tf.array_layers = 1
	tf.mipmaps = 1
	tf.usage_bits = (
		RenderingDevice.TEXTURE_USAGE_SAMPLING_BIT
		| RenderingDevice.TEXTURE_USAGE_COLOR_ATTACHMENT_BIT
		| RenderingDevice.TEXTURE_USAGE_STORAGE_BIT
		| RenderingDevice.TEXTURE_USAGE_CAN_UPDATE_BIT
		| RenderingDevice.TEXTURE_USAGE_CAN_COPY_TO_BIT
	)


	var agent_uniforms = []
	for i in 2:
		texture_rds[i] = rd.texture_create(tf, RDTextureView.new(), [])
		# Make sure our textures are cleared.
		rd.texture_clear(texture_rds[i], Color.BLACK, 0, 1, 0, 1)

		var trail_uniform := _create_trail_uniform(texture_rds[i])
		texture_sets[i] = rd.uniform_set_create([trail_uniform], decay_shader, 0)
		agent_uniform_sets[i] = rd.uniform_set_create([trail_uniform], agent_shader, 0)

	
	var uniform := RDUniform.new()
	uniform.uniform_type = RenderingDevice.UNIFORM_TYPE_STORAGE_BUFFER
	uniform.binding = 1
	uniform.add_id(agent_buffer)
	agent_set = rd.uniform_set_create([uniform], agent_shader, 0)


func _render_process(with_next_texture: int, tex_size: Vector2i, decay_fctr: float) -> void:
	var next_set := texture_sets[with_next_texture]
	var current_set := texture_sets[(with_next_texture - 1) % 2]

	# Run agents
	var push_constant := PackedFloat32Array()
	push_constant.push_back(tex_size.x)
	push_constant.push_back(tex_size.y)
	push_constant.push_back(0.0)
	push_constant.push_back(0.0)


	var agent_compute_list := rd.compute_list_begin()
	rd.compute_list_bind_compute_pipeline(agent_compute_list, agent_pipeline)
	rd.compute_list_bind_uniform_set(agent_compute_list, current_set, 0)
	rd.compute_list_bind_uniform_set(agent_compute_list, next_set, 1)
	rd.compute_list_bind_uniform_set(agent_compute_list, agent_set, 2)
	rd.compute_list_set_push_constant(
		agent_compute_list, push_constant.to_byte_array(), push_constant.size() * 4
	)
	rd.compute_list_dispatch(agent_compute_list, 16, 1, 1)
	rd.compute_list_end()

	# Decay and dissipate
	push_constant.clear()
	push_constant.push_back(tex_size.x)
	push_constant.push_back(tex_size.y)
	push_constant.push_back(decay_fctr)
	push_constant.push_back(0.0)

	@warning_ignore("integer_division")
	var x_groups := (tex_size.x - 1) / 8 + 1
	@warning_ignore("integer_division")
	var y_groups := (tex_size.y - 1) / 8 + 1

	# Run the compute shader
	var compute_list := rd.compute_list_begin()
	rd.compute_list_bind_compute_pipeline(compute_list, decay_pipeline)
	rd.compute_list_bind_uniform_set(compute_list, current_set, 0)
	rd.compute_list_bind_uniform_set(compute_list, next_set, 1)
	rd.compute_list_set_push_constant(
		compute_list, push_constant.to_byte_array(), push_constant.size() * 4
	)
	rd.compute_list_dispatch(compute_list, x_groups, y_groups, 1)
	rd.compute_list_end()

	# We don't need to sync up here, Godots default barriers will do the trick.
	# If you want the output of a compute shader to be used as input of
	# another computer shader you'll need to add a barrier:
	#rd.barrier(RenderingDevice.BARRIER_MASK_COMPUTE)


func _free_compute_resources() -> void:
	# Note that our sets and pipeline are cleaned up automatically
	# as they are dependencies :P
	for i in 3:
		if texture_rds[i]:
			rd.free_rid(texture_rds[i])

	if decay_shader:
		rd.free_rid(decay_shader)