internal/shader: restrict the number of arguments for atan/atan2

xiaoshengduan · Mar 26, 2022 · a936ffc · a936ffc
1 parent d30a31b
commit a936ffc
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 1 deletion.
diff --git 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
@@ -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)
+	}
+}