diff --git a/src/ast/AST.go b/src/ast/AST.go
index 0950b2f..87daccf 100644
--- a/src/ast/AST.go
+++ b/src/ast/AST.go
@@ -1,6 +1,9 @@
 package ast
 
-import "git.urbach.dev/cli/q/src/expression"
+import (
+	"git.urbach.dev/cli/q/src/expression"
+	"git.urbach.dev/cli/q/src/token"
+)
 
 type (
 	Node any
@@ -34,6 +37,7 @@ type (
 		Body AST
 	}
 	Return struct {
+		Token  token.Token
 		Values []*expression.Expression
 	}
 	Switch struct {
diff --git a/src/ast/parseKeyword.go b/src/ast/parseKeyword.go
index 7aa9ffe..74a28fa 100644
--- a/src/ast/parseKeyword.go
+++ b/src/ast/parseKeyword.go
@@ -59,11 +59,11 @@ func parseKeyword(tokens token.List, file *fs.File, nodes AST) (Node, error) {
 
 	case token.Return:
 		if len(tokens) == 1 {
-			return &Return{}, nil
+			return &Return{Token: tokens[0]}, nil
 		}
 
 		values := expression.NewList(tokens[1:])
-		return &Return{Values: values}, nil
+		return &Return{Token: tokens[0], Values: values}, nil
 
 	case token.Switch:
 		blockStart := tokens.IndexKind(token.BlockStart)
diff --git a/src/core/Compile.go b/src/core/Compile.go
index dded436..0474050 100644
--- a/src/core/Compile.go
+++ b/src/core/Compile.go
@@ -2,7 +2,6 @@ package core
 
 import (
 	"git.urbach.dev/cli/q/src/ast"
-	"git.urbach.dev/cli/q/src/ssa"
 )
 
 // Compile translates tokens to SSA form.
@@ -22,9 +21,5 @@ func (f *Function) Compile() {
 		return
 	}
 
-	if f.needsReturn() {
-		f.Block().Append(&ssa.Return{})
-	}
-
 	f.Err = f.optimize()
 }
\ No newline at end of file
diff --git a/src/core/compileReturn.go b/src/core/compileReturn.go
index 0fa1f10..69669d1 100644
--- a/src/core/compileReturn.go
+++ b/src/core/compileReturn.go
@@ -8,15 +8,21 @@ import (
 
 // compileReturn compiles a return instruction.
 func (f *Function) compileReturn(node *ast.Return) error {
+	if len(node.Values) != len(f.Output) {
+		position := node.Token.End()
+
+		if len(node.Values) > 0 {
+			position = node.Values[0].Token.Position
+		}
+
+		return errors.New(&ReturnCountMismatch{Count: len(node.Values), ExpectedCount: len(f.Output)}, f.File, position)
+	}
+
 	if len(node.Values) == 0 {
 		f.Append(&ssa.Return{})
 		return nil
 	}
 
-	if len(node.Values) != len(f.Output) {
-		return errors.New(&ReturnCountMismatch{Count: len(node.Values), ExpectedCount: len(f.Output)}, f.File, node.Values[0].Token.Position)
-	}
-
 	args, err := f.decompose(node.Values, f.Output, true)
 
 	if err != nil {
diff --git a/src/core/errors_test.go b/src/core/errors_test.go
index e160b95..bd6c4fe 100644
--- a/src/core/errors_test.go
+++ b/src/core/errors_test.go
@@ -48,6 +48,8 @@ var errs = []struct {
 	{"ResourceTypeMismatch.q", &core.TypeMismatch{Encountered: "int64", Expected: "!int64", ParameterName: "_", IsReturn: false}},
 	{"ReturnCountMismatch.q", &core.ReturnCountMismatch{Count: 1, ExpectedCount: 0}},
 	{"ReturnCountMismatch2.q", &core.ReturnCountMismatch{Count: 1, ExpectedCount: 2}},
+	{"ReturnCountMismatch3.q", &core.ReturnCountMismatch{Count: 0, ExpectedCount: 1}},
+	{"ReturnCountMismatch4.q", &core.ReturnCountMismatch{Count: 0, ExpectedCount: 1}},
 	{"TypeMismatch.q", &core.TypeMismatch{Encountered: "string", Expected: "int64", ParameterName: "x", IsReturn: false}},
 	{"TypeMismatch2.q", &core.TypeMismatch{Encountered: "string", Expected: "int64", ParameterName: "y", IsReturn: true}},
 	{"UndefinedStructField.q", &core.UndefinedStructField{Identifier: "p", FieldName: "y"}},
diff --git a/src/core/needsReturn.go b/src/core/needsReturn.go
index 7b10374..fd7acad 100644
--- a/src/core/needsReturn.go
+++ b/src/core/needsReturn.go
@@ -8,6 +8,12 @@ func (f *Function) needsReturn() bool {
 		return false
 	}
 
-	_, endsWithReturn := f.Block().Last().(*ssa.Return)
+	lastBlock := f.Block()
+
+	if lastBlock.Loop != nil {
+		return false
+	}
+
+	_, endsWithReturn := lastBlock.Last().(*ssa.Return)
 	return !endsWithReturn
 }
\ No newline at end of file
diff --git a/src/core/optimize.go b/src/core/optimize.go
index 36f0b41..293d930 100644
--- a/src/core/optimize.go
+++ b/src/core/optimize.go
@@ -1,6 +1,7 @@
 package core
 
 import (
+	"git.urbach.dev/cli/q/src/errors"
 	"git.urbach.dev/cli/q/src/fold"
 	"git.urbach.dev/cli/q/src/ssa"
 )
@@ -12,6 +13,17 @@ func (f *Function) optimize() error {
 	// unnecessary return statements in the later phases.
 	f.removeDeadBlocks()
 
+	// After the removal of dead blocks, if the last block is
+	// not part of a loop and did not end with a return
+	// statement, an implicit return is inserted.
+	if f.needsReturn() {
+		if len(f.Output) > 0 {
+			return errors.New(&ReturnCountMismatch{Count: 0, ExpectedCount: len(f.Output)}, f.File, f.Output[0].StartPos)
+		}
+
+		f.Block().Append(&ssa.Return{})
+	}
+
 	// Binary operations with constant operands are evaluated
 	// at compile time. For example, 1 + 2 becomes 3, and the
 	// result is propagated to subsequent operations.
diff --git a/src/core/testdata/ReturnCountMismatch3.q b/src/core/testdata/ReturnCountMismatch3.q
new file mode 100644
index 0000000..7469550
--- /dev/null
+++ b/src/core/testdata/ReturnCountMismatch3.q
@@ -0,0 +1,7 @@
+main() {
+	f()
+}
+
+f() -> int {
+	return
+}
\ No newline at end of file
diff --git a/src/core/testdata/ReturnCountMismatch4.q b/src/core/testdata/ReturnCountMismatch4.q
new file mode 100644
index 0000000..ae957e3
--- /dev/null
+++ b/src/core/testdata/ReturnCountMismatch4.q
@@ -0,0 +1,5 @@
+main() {
+	f()
+}
+
+f() -> int {}
\ No newline at end of file
diff --git a/tests/shift-register.q b/tests/shift-register.q
index dce24c5..79081ac 100644
--- a/tests/shift-register.q
+++ b/tests/shift-register.q
@@ -2,6 +2,6 @@ main() {
 	f(1, 2)
 }
 
-f(x int, y int) -> int {
+f(x int, y int) {
 	assert x << y == 0b100
 }
\ No newline at end of file