fix: force adding a return at end of function

This will avoid infinite loops, and is necessary for functions with
no returned values. This doesn't remove the need to better check the
consistency of return statements in general.

2 files changed, 10 insertions, 0 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..c13c33d 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([]scanner.Token{{Id: lang.Return}})
+		if err != nil {
+			return out, err
+		}
+		out = append(out, x...)
+	}
 	out = append(out, scanner.Token{Id: lang.Label, Str: fname + "_end"})
 	return out, err
 }