4 files changed, 11 insertions, 3 deletions
diff --git a/pkg/lang/compiler/compiler.go b/pkg/lang/compiler/compiler.go
index ae2a71c..728dec1 100644
--- a/pkg/lang/compiler/compiler.go
+++ b/pkg/lang/compiler/compiler.go
@@ -50,7 +50,7 @@ func (comp *Compiler) Compile() (code.Code, error) {
 }
 
 func (comp *Compiler) preDeclareFunction(fnDeclStmt ast.StmtFnDecl) error {
-	if _, ok := comp.scopes.DeclareFunction(fnDeclStmt.Name.Value); !ok {
+	if _, ok := comp.scopes.DeclareFunction(fnDeclStmt.Name.Value, uint(len(fnDeclStmt.Args))); !ok {
 		return fmt.Errorf("function %s already declared", fnDeclStmt.Name.Value)
 	}
 
@@ -103,7 +103,7 @@ func (comp *Compiler) compileStmt(t *code.Builder, stmt ast.Stmt) error {
 }
 
 func (comp *Compiler) compileFnDeclStmt(t *code.Builder, fnDeclStmt ast.StmtFnDecl) error {
-	marker, ok := comp.scopes.DeclareFunction(fnDeclStmt.Name.Value)
+	marker, ok := comp.scopes.DeclareFunction(fnDeclStmt.Name.Value, uint(len(fnDeclStmt.Args)))
 	if !ok {
 		// If we are in the root scope, the function was simply predeclared :)
 		if comp.scopes.IsRootScope() {
@@ -598,6 +598,11 @@ func (comp *Compiler) compileIdentExpr(t *code.Builder, expr ast.ExprIdent) erro
 
 		t.AppendOp(code.OpPushFunction)
 		t.AppendMarkerReference(symbol.data.marker)
+
+		if symbol.data.args != 0 {
+			t.AppendOp(code.OpSetArgCount)
+			t.AppendInt(int64(symbol.data.args))
+		}
 	default:
 		panic(fmt.Errorf("unknown symbol kind: %v", symbolId.symbolKind))
 	}
diff --git a/pkg/lang/compiler/compiler_test.go b/pkg/lang/compiler/compiler_test.go
index cd62088..bdbc375 100644
--- a/pkg/lang/compiler/compiler_test.go
+++ b/pkg/lang/compiler/compiler_test.go
@@ -383,6 +383,7 @@ func TestFunctionArgs(t *testing.T) {
 
 	expected := `
 	push_function @add
+	set_arg_count 2
 	push_int 4
 	push_int 5
 	call 2
diff --git a/pkg/lang/compiler/scope_chain.go b/pkg/lang/compiler/scope_chain.go
index 3a3b819..6705852 100644
--- a/pkg/lang/compiler/scope_chain.go
+++ b/pkg/lang/compiler/scope_chain.go
@@ -91,7 +91,7 @@ func (sc *ScopeChain) Declare(name string) (int, bool) {
 	return indexInScope, true
 }
 
-func (sc *ScopeChain) DeclareFunction(name string) (code.Marker, bool) {
+func (sc *ScopeChain) DeclareFunction(name string, args uint) (code.Marker, bool) {
 	if _, ok := sc.nameToSymbol[name]; ok {
 		return "", false
 	}
@@ -111,6 +111,7 @@ func (sc *ScopeChain) DeclareFunction(name string) (code.Marker, bool) {
 		name: name,
 		data: SymbolFunction{
 			marker: unitName,
+			args:   args,
 		},
 	})
 
diff --git a/pkg/lang/compiler/symbol.go b/pkg/lang/compiler/symbol.go
index d22cdc0..453cb2b 100644
--- a/pkg/lang/compiler/symbol.go
+++ b/pkg/lang/compiler/symbol.go
@@ -35,4 +35,5 @@ type SymbolVariable struct {
 
 type SymbolFunction struct {
 	marker code.Marker
+	args   uint
 }