From a936ffc032890a6a556c8728bb42a4d13a51f769 Mon Sep 17 00:00:00 2001
From: Hajime Hoshi <hajimehoshi@gmail.com>
Date: Sat, 26 Mar 2022 19:39:59 +0900
Subject: [PATCH] internal/shader: restrict the number of arguments for
 atan/atan2

---
 internal/shader/expr.go        | 11 +++++++++-
 internal/shader/syntax_test.go | 37 ++++++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/internal/shader/expr.go b/internal/shader/expr.go
index 7ba1fd4f9..c03370f23 100644
--- a/internal/shader/expr.go
+++ b/internal/shader/expr.go
@@ -365,8 +365,17 @@ func (cs *compileState) parseExpr(block *block, expr ast.Expr, markLocalVariable
 						},
 					}, []shaderir.Type{{Main: shaderir.Float}}, stmts, true
 				}
+			case shaderir.Atan:
+				if len(args) != 1 {
+					cs.addError(e.Pos(), fmt.Sprintf("number of %s's arguments must be 1 but %d", callee.BuiltinFunc, len(args)))
+					return nil, nil, nil, false
+				}
+			case shaderir.Atan2:
+				if len(args) != 2 {
+					cs.addError(e.Pos(), fmt.Sprintf("number of %s's arguments must be 2 but %d", callee.BuiltinFunc, len(args)))
+					return nil, nil, nil, false
+				}
 			}
-
 			var t shaderir.Type
 			switch callee.BuiltinFunc {
 			case shaderir.BoolF:
diff --git a/internal/shader/syntax_test.go b/internal/shader/syntax_test.go
index 69a11bad2..9f2df3db3 100644
--- a/internal/shader/syntax_test.go
+++ b/internal/shader/syntax_test.go
@@ -1146,3 +1146,40 @@ func Fragment(position vec4, texCoord vec2, color vec4) vec4 {
 		}
 	}
 }
+
+func TestSyntaxAtan(t *testing.T) {
+	// `atan` takes 1 argument.
+	if _, err := compileToIR([]byte(`package main
+
+func Fragment(position vec4, texCoord vec2, color vec4) vec4 {
+	return atan(vec4(0))
+}
+`)); err != nil {
+		t.Error(err)
+	}
+	if _, err := compileToIR([]byte(`package main
+
+func Fragment(position vec4, texCoord vec2, color vec4) vec4 {
+	return atan(vec4(0), vec4(0))
+}
+`)); err == nil {
+		t.Errorf("error must be non-nil but was nil")
+	}
+	// `atan2` takes 2 arguments.
+	if _, err := compileToIR([]byte(`package main
+
+func Fragment(position vec4, texCoord vec2, color vec4) vec4 {
+	return atan2(vec4(0))
+}
+`)); err == nil {
+		t.Errorf("error must be non-nil but was nil")
+	}
+	if _, err := compileToIR([]byte(`package main
+
+func Fragment(position vec4, texCoord vec2, color vec4) vec4 {
+	return atan2(vec4(0), vec4(0))
+}
+`)); err != nil {
+		t.Error(err)
+	}
+}