Merge pull request #4 from mvertes/fix-no-return-func

fix: force a return at end of function

2 files changed, 13 insertions, 1 deletions

diff --git a/parser/interpreter_test.go b/parser/interpreter_test.go
index c3f2909..8f166df 100644
--- a/parser/interpreter_test.go
+++ b/parser/interpreter_test.go
@@ -97,6 +97,7 @@ func TestFunc(t *testing.T) {
 		{src: "func f(a int) int {return a+2}; f(5) - f(3)", res: "2"},
 		{src: "func f(a int) int {return a+2}; f(3) - 2", res: "3"},
 		{src: "func f(a, b, c int) int {return a+b-c} ; f(7, 1, 3)", res: "5"},
+		{src: "var a int; func f() {a = a+2}; f(); a", res: "2"},
 	})
 }
 
diff --git a/parser/parse.go b/parser/parse.go
index ffa8102..547957b 100644
--- a/parser/parse.go
+++ b/parser/parse.go
@@ -246,6 +246,15 @@ func (p *Parser) ParseFunc(in Tokens) (out Tokens, err error) {
 		out = append(out, scanner.Token{Id: lang.Grow, Beg: l})
 	}
 	out = append(out, toks...)
+	if out[len(out)-1].Id != lang.Return {
+		// Ensure that a return statment is always added at end of function.
+		// TODO: detect missing or wrong returns.
+		x, err := p.ParseReturn(nil)
+		if err != nil {
+			return out, err
+		}
+		out = append(out, x...)
+	}
 	out = append(out, scanner.Token{Id: lang.Label, Str: fname + "_end"})
 	return out, err
 }
@@ -411,10 +420,12 @@ func (p *Parser) ParseLabel(in Tokens) (out Tokens, err error) {
 }
 
 func (p *Parser) ParseReturn(in Tokens) (out Tokens, err error) {
-	if len(in) > 1 {
+	if l := len(in); l > 1 {
 		if out, err = p.ParseExpr(in[1:]); err != nil {
 			return out, err
 		}
+	} else if l == 0 {
+		in = Tokens{{Id: lang.Return}} // Implicit return in functions with no return parameters.
 	}
 
 	// TODO: the function symbol should be already present in the parser context.