From 9693ce83821eb10ea434342dfe812435d2dfaf6f Mon Sep 17 00:00:00 2001
From: Hajime Hoshi <hajimehoshi@gmail.com>
Date: Sun, 25 Aug 2024 18:57:33 +0900
Subject: [PATCH] internal/shader: allow less arguments at Fragment

Closes #3073
---
 internal/builtinshader/shader.go |   4 +-
 internal/shader/shader.go        |   6 ++
 shader_test.go                   | 103 +++++++++++++++++++++++++++++++
 3 files changed, 111 insertions(+), 2 deletions(-)

diff --git a/internal/builtinshader/shader.go b/internal/builtinshader/shader.go
index 3e168f2e5..b543070d8 100644
--- a/internal/builtinshader/shader.go
+++ b/internal/builtinshader/shader.go
@@ -170,7 +170,7 @@ var ScreenShaderSource = []byte(`//kage:unit pixels
 
 package main
 
-func Fragment(dstPos vec4, srcPos vec2, color vec4) vec4 {
+func Fragment(dstPos vec4, srcPos vec2) vec4 {
 	// Blend source colors in a square region, which size is 1/scale.
 	scale := imageDstSize()/imageSrc0Size()
 	pos := srcPos
@@ -193,7 +193,7 @@ var ClearShaderSource = []byte(`//kage:unit pixels
 
 package main
 
-func Fragment(dstPos vec4, srcPos vec2, color vec4) vec4 {
+func Fragment() vec4 {
 	return vec4(0)
 }
 `)
diff --git a/internal/shader/shader.go b/internal/shader/shader.go
index 1aa6e7b1b..3952383c0 100644
--- a/internal/shader/shader.go
+++ b/internal/shader/shader.go
@@ -821,6 +821,12 @@ func (cs *compileState) parseFunc(block *block, d *ast.FuncDecl) (function, bool
 
 	inParams, outParams, returnType := cs.parseFuncParams(block, d.Name.Name, d)
 	if d.Name.Name == cs.fragmentEntry {
+		if len(inParams) == 0 {
+			inParams = append(inParams, variable{
+				name: "_",
+				typ:  shaderir.Type{Main: shaderir.Vec4},
+			})
+		}
 		// The 0th inParams is a special variable for position and is not included in varying variables.
 		if diff := len(cs.ir.Varyings) - (len(inParams) - 1); diff > 0 {
 			// inParams is not enough when the vertex shader has more returning values than the fragment shader's arguments.
diff --git a/shader_test.go b/shader_test.go
index 70c292524..910f5de04 100644
--- a/shader_test.go
+++ b/shader_test.go
@@ -2682,3 +2682,106 @@ func Fragment(dstPos vec4, srcPos vec2, color vec4, custom vec4) vec4 {
 		}
 	}
 }
+
+func TestShaderFragmentLessArguments(t *testing.T) {
+	const w, h = 16, 16
+
+	s0, err := ebiten.NewShader([]byte(`//kage:unit pixels
+
+package main
+
+func Fragment() vec4 {
+	return vec4(1, 0, 0, 1)
+}
+`))
+	if err != nil {
+		t.Fatal(err)
+	}
+	s1, err := ebiten.NewShader([]byte(`//kage:unit pixels
+
+package main
+
+func Fragment(dstPos vec4) vec4 {
+	return vec4(0, 1, 0, 1)
+}
+`))
+	if err != nil {
+		t.Fatal(err)
+	}
+	s2, err := ebiten.NewShader([]byte(`//kage:unit pixels
+
+package main
+
+func Fragment(dstPos vec4, srcPos vec2) vec4 {
+	return vec4(0, 0, 1, 1)
+}
+`))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	dst := ebiten.NewImage(w, h)
+	for idx, s := range []*ebiten.Shader{s0, s1, s2} {
+		dst.Clear()
+		dst.DrawTrianglesShader([]ebiten.Vertex{
+			{
+				DstX:   0,
+				DstY:   0,
+				SrcX:   0,
+				SrcY:   0,
+				ColorR: 1,
+				ColorG: 1,
+				ColorB: 1,
+				ColorA: 1,
+			},
+			{
+				DstX:   w,
+				DstY:   0,
+				SrcX:   w,
+				SrcY:   0,
+				ColorR: 1,
+				ColorG: 1,
+				ColorB: 1,
+				ColorA: 1,
+			},
+			{
+				DstX:   0,
+				DstY:   h,
+				SrcX:   0,
+				SrcY:   h,
+				ColorR: 1,
+				ColorG: 1,
+				ColorB: 1,
+				ColorA: 1,
+			},
+			{
+				DstX:   w,
+				DstY:   h,
+				SrcX:   w,
+				SrcY:   h,
+				ColorR: 1,
+				ColorG: 1,
+				ColorB: 1,
+				ColorA: 1,
+			},
+		}, []uint16{0, 1, 2, 1, 2, 3}, s, nil)
+
+		for j := 0; j < h; j++ {
+			for i := 0; i < w; i++ {
+				got := dst.At(i, j).(color.RGBA)
+				var want color.RGBA
+				switch idx {
+				case 0:
+					want = color.RGBA{R: 0xff, A: 0xff}
+				case 1:
+					want = color.RGBA{G: 0xff, A: 0xff}
+				case 2:
+					want = color.RGBA{B: 0xff, A: 0xff}
+				}
+				if !sameColors(got, want, 2) {
+					t.Errorf("dst.At(%d, %d): got: %v, want: %v", i, j, got, want)
+				}
+			}
+		}
+	}
+}