The return statement now accepts any expression as its return value

So something like 'return ++user_x;' is now possible
Admittedly this needed quite a bit of refactoring mainly due to the fact that return types now have to be checked after resolving the function rather than before

src/thingdef/thingdef.cpp

@@ -270,23 +270,28 @@ static void FinishThingdef()
 	for (i = 0; i < StateTempCalls.Size(); ++i)
 	{
 		FStateTempCall *tcall = StateTempCalls[i];
-		VMFunction *func;
+		VMFunction *func = nullptr;
 
 		assert(tcall->Code != NULL);
 
+		// We don't know the return type in advance for anonymous functions.
+		FCompileContext ctx(tcall->ActorClass, nullptr);
+		tcall->Code = tcall->Code->Resolve(ctx);
+		tcall->Proto = ctx.ReturnProto;
+
+		// Make sure resolving it didn't obliterate it.
+		if (tcall->Code != nullptr)
+		{
 			// Can we call this function directly without wrapping it in an
 			// anonymous function? e.g. Are we passing any parameters to it?
 			func = tcall->Code->GetDirectFunction();
-		if (func == NULL)
-		{
-			FCompileContext ctx(tcall->ActorClass);
-			tcall->Code = tcall->Code->Resolve(ctx);
 
-			// Make sure resolving it didn't obliterate it.
+			if (func == nullptr)
-			if (tcall->Code != NULL)
 			{
 				VMFunctionBuilder buildit;
 
+				assert(tcall->Proto != nullptr);
+
 				// Allocate registers used to pass parameters in.
 				// self, stateowner, state (all are pointers)
 				buildit.Registers[REGT_POINTER].Get(3);
@@ -300,15 +305,7 @@ static void FinishThingdef()
 				// Generate prototype for this anonymous function
 				TArray<PType *> args(3);
 				SetImplicitArgs(&args, NULL, tcall->ActorClass, VARF_Method | VARF_Action);
-				if (tcall->Proto != NULL)
-				{
 				sfunc->Proto = NewPrototype(tcall->Proto->ReturnTypes, args);
-				}
-				else
-				{
-					TArray<PType *> norets(0);
-					sfunc->Proto = NewPrototype(norets, args);
-				}
 
 				func = sfunc;
 
@@ -321,11 +318,9 @@ static void FinishThingdef()
 					codesize += sfunc->CodeSize;
 				}
 			}
-		}
+
-		if (tcall->Code != NULL)
-		{
 			delete tcall->Code;
-			tcall->Code = NULL;
+			tcall->Code = nullptr;
 			for (int k = 0; k < tcall->NumStates; ++k)
 			{
 				tcall->ActorClass->OwnedStates[tcall->FirstState + k].SetAction(func);

src/thingdef/thingdef.h

@@ -190,10 +190,9 @@ AFuncDesc *FindFunction(const char * string);
 void ParseStates(FScanner &sc, PClassActor *actor, AActor *defaults, Baggage &bag);
 void ParseFunctionParameters(FScanner &sc, PClassActor *cls, TArray<FxExpression *> &out_params,
 	PFunction *afd, FString statestring, FStateDefinitions *statedef);
-FxExpression *ParseActions(FScanner &sc, FState state, FString statestring, Baggage &bag, PPrototype *&proto, bool &endswithret);
+FxExpression *ParseActions(FScanner &sc, FState state, FString statestring, Baggage &bag, bool &endswithret);
 class FxVMFunctionCall *ParseAction(FScanner &sc, FState state, FString statestring, Baggage &bag);
 FName CheckCastKludges(FName in);
-void AddImplicitReturn(class FxSequence *code, const PPrototype *proto, FScanner &sc);
 void SetImplicitArgs(TArray<PType *> *args, TArray<DWORD> *argflags, PClassActor *cls, DWORD funcflags);
 
 PFunction *FindGlobalActionFunction(const char *name);
@@ -356,6 +355,7 @@ int MatchString (const char *in, const char **strings);
 
 
 #define ACTION_RETURN_STATE(v) do { FState *state = v; if (numret > 0) { assert(ret != NULL); ret->SetPointer(state, ATAG_STATE); return 1; } return 0; } while(0)
+#define ACTION_RETURN_FLOAT(v) do { double u = v; if (numret > 0) { assert(ret != nullptr); ret->SetFloat(u); return 1; } return 0; } while(0)
 #define ACTION_RETURN_INT(v) do { int u = v; if (numret > 0) { assert(ret != NULL); ret->SetInt(u); return 1; } return 0; } while(0)
 #define ACTION_RETURN_BOOL(v) ACTION_RETURN_INT(v)
 

src/thingdef/thingdef_exp.h

@@ -62,14 +62,16 @@ class FxJumpStatement;
 struct FCompileContext
 {
 	TArray<FxJumpStatement *> Jumps;
+	PPrototype *ReturnProto;
 	PClassActor *Class;
 
-	FCompileContext(PClassActor *cls = nullptr);
+	FCompileContext(PClassActor *cls = nullptr, PPrototype *ret = nullptr);
 
 	PSymbol *FindInClass(FName identifier);
 	PSymbol *FindGlobal(FName identifier);
 
 	void HandleJumps(int token, FxExpression *handler);
+	void CheckReturn(PPrototype *proto, FScriptPosition &pos);
 };
 
 //==========================================================================
@@ -170,14 +172,14 @@ struct ExpVal
 
 struct ExpEmit
 {
-	ExpEmit() : RegNum(0), RegType(REGT_NIL), Konst(false), Fixed(false) {}
+	ExpEmit() : RegNum(0), RegType(REGT_NIL), Konst(false), Fixed(false), Final(false) {}
-	ExpEmit(int reg, int type) : RegNum(reg), RegType(type), Konst(false), Fixed(false) {}
+	ExpEmit(int reg, int type) : RegNum(reg), RegType(type), Konst(false), Fixed(false), Final(false) {}
-	ExpEmit(int reg, int type, bool konst)  : RegNum(reg), RegType(type), Konst(konst), Fixed(false) {}
+	ExpEmit(int reg, int type, bool konst)  : RegNum(reg), RegType(type), Konst(konst), Fixed(false), Final(false) {}
 	ExpEmit(VMFunctionBuilder *build, int type);
 	void Free(VMFunctionBuilder *build);
 	void Reuse(VMFunctionBuilder *build);
 
-	BYTE RegNum, RegType, Konst:1, Fixed:1;
+	BYTE RegNum, RegType, Konst:1, Fixed:1, Final:1;
 };
 
 //==========================================================================
@@ -202,6 +204,7 @@ public:
 	
 	virtual bool isConstant() const;
 	virtual bool RequestAddress();
+	virtual PPrototype *ReturnProto();
 	virtual VMFunction *GetDirectFunction();
 	bool IsNumeric() const { return ValueType != TypeName && (ValueType->GetRegType() == REGT_INT || ValueType->GetRegType() == REGT_FLOAT); }
 	bool IsPointer() const { return ValueType->GetRegType() == REGT_POINTER; }
@@ -748,6 +751,7 @@ public:
 class FxRandom : public FxExpression
 {
 protected:
+	bool EmitTail;
 	FRandom *rng;
 	FxExpression *min, *max;
 
@@ -756,7 +760,7 @@ public:
 	FxRandom(FRandom *, FxExpression *mi, FxExpression *ma, const FScriptPosition &pos);
 	~FxRandom();
 	FxExpression *Resolve(FCompileContext&);
-
+	PPrototype *ReturnProto();
 	ExpEmit Emit(VMFunctionBuilder *build);
 };
 
@@ -802,6 +806,7 @@ public:
 
 class FxRandom2 : public FxExpression
 {
+	bool EmitTail;
 	FRandom * rng;
 	FxExpression *mask;
 
@@ -810,7 +815,7 @@ public:
 	FxRandom2(FRandom *, FxExpression *m, const FScriptPosition &pos);
 	~FxRandom2();
 	FxExpression *Resolve(FCompileContext&);
-
+	PPrototype *ReturnProto();
 	ExpEmit Emit(VMFunctionBuilder *build);
 };
 
@@ -915,8 +920,9 @@ public:
 
 class FxActionSpecialCall : public FxExpression
 {
-	FxExpression *Self;
 	int Special;
+	bool EmitTail;
+	FxExpression *Self;
 	FArgumentList *ArgList;
 
 public:
@@ -924,6 +930,7 @@ public:
 	FxActionSpecialCall(FxExpression *self, int special, FArgumentList *args, const FScriptPosition &pos);
 	~FxActionSpecialCall();
 	FxExpression *Resolve(FCompileContext&);
+	PPrototype *ReturnProto();
 	ExpEmit Emit(VMFunctionBuilder *build);
 };
 
@@ -954,6 +961,7 @@ public:
 
 class FxVMFunctionCall : public FxExpression
 {
+	bool EmitTail;
 	PFunction *Function;
 	FArgumentList *ArgList;
 
@@ -961,13 +969,10 @@ public:
 	FxVMFunctionCall(PFunction *func, FArgumentList *args, const FScriptPosition &pos);
 	~FxVMFunctionCall();
 	FxExpression *Resolve(FCompileContext&);
+	PPrototype *ReturnProto();
+	VMFunction *GetDirectFunction();
 	ExpEmit Emit(VMFunctionBuilder *build);
-	ExpEmit Emit(VMFunctionBuilder *build, bool tailcall);
 	bool CheckEmitCast(VMFunctionBuilder *build, bool returnit, ExpEmit &reg);
-	unsigned GetArgCount() const { return ArgList == NULL ? 0 : ArgList->Size(); }
-	PFunction *GetFunction() const { return Function; }
-	VMFunction *GetVMFunction() const { return Function->Variants[0].Implementation; }
-	bool IsDirectFunction();
 };
 
 //==========================================================================
@@ -1089,10 +1094,10 @@ public:
 
 class FxReturnStatement : public FxExpression
 {
-	FxVMFunctionCall *Call;
+	FxExpression *Value;
 
 public:
-	FxReturnStatement(FxVMFunctionCall *call, const FScriptPosition &pos);
+	FxReturnStatement(FxExpression *value, const FScriptPosition &pos);
 	~FxReturnStatement();
 	FxExpression *Resolve(FCompileContext&);
 	ExpEmit Emit(VMFunctionBuilder *build);
@@ -1146,12 +1151,14 @@ public:
 
 class FxRuntimeStateIndex : public FxExpression
 {
+	bool EmitTail;
 	FxExpression *Index;
 
 public:
 	FxRuntimeStateIndex(FxExpression *index);
 	~FxRuntimeStateIndex();
 	FxExpression *Resolve(FCompileContext&);
+	PPrototype *ReturnProto();
 	ExpEmit Emit(VMFunctionBuilder *build);
 };
 

src/thingdef/thingdef_expression.cpp

@@ -91,7 +91,7 @@ static const FLOP FxFlops[] =
 //
 //==========================================================================
 
-FCompileContext::FCompileContext(PClassActor *cls) : Class(cls)
+FCompileContext::FCompileContext(PClassActor *cls, PPrototype *ret) : Class(cls), ReturnProto(ret)
 {
 }
 
@@ -118,6 +118,50 @@ void FCompileContext::HandleJumps(int token, FxExpression *handler)
 	}
 }
 
+void FCompileContext::CheckReturn(PPrototype *proto, FScriptPosition &pos)
+{
+	assert(proto != nullptr);
+	bool fail = false;
+
+	if (ReturnProto == nullptr)
+	{
+		ReturnProto = proto;
+		return;
+	}
+
+	// A prototype that defines fewer return types can be compatible with
+	// one that defines more if the shorter one matches the initial types
+	// for the longer one.
+	if (ReturnProto->ReturnTypes.Size() < proto->ReturnTypes.Size())
+	{ // Make proto the shorter one to avoid code duplication below.
+		swapvalues(proto, ReturnProto);
+	}
+	// If one prototype returns nothing, they both must.
+	if (proto->ReturnTypes.Size() == 0)
+	{
+		if (ReturnProto->ReturnTypes.Size() != 0)
+		{
+			fail = true;
+		}
+	}
+	else
+	{
+		for (unsigned i = 0; i < proto->ReturnTypes.Size(); i++)
+		{
+			if (ReturnProto->ReturnTypes[i] != proto->ReturnTypes[i])
+			{ // Incompatible
+				fail = true;
+				break;
+			}
+		}
+	}
+
+	if (fail)
+	{
+		pos.Message(MSG_ERROR, "All return expressions must deduce to the same type");
+	}
+}
+
 //==========================================================================
 //
 // ExpEmit
@@ -125,7 +169,7 @@ void FCompileContext::HandleJumps(int token, FxExpression *handler)
 //==========================================================================
 
 ExpEmit::ExpEmit(VMFunctionBuilder *build, int type)
-: RegNum(build->Registers[type].Get(1)), RegType(type), Konst(false), Fixed(false)
+: RegNum(build->Registers[type].Get(1)), RegType(type), Konst(false), Fixed(false), Final(false)
 {
 }
 
@@ -228,6 +272,26 @@ bool FxExpression::RequestAddress()
 	return false;
 }
 
+//==========================================================================
+//
+// Called by return statements.
+//
+//==========================================================================
+
+PPrototype *FxExpression::ReturnProto()
+{
+	assert(ValueType != nullptr);
+
+	TArray<PType *> ret(0);
+	TArray<PType *> none(0);
+	if (ValueType != TypeVoid)
+	{
+		ret.Push(ValueType);
+	}
+
+	return NewPrototype(ret, none);
+}
+
 //==========================================================================
 //
 //
@@ -2651,6 +2715,7 @@ ExpEmit FxMinMax::Emit(VMFunctionBuilder *build)
 FxRandom::FxRandom(FRandom * r, FxExpression *mi, FxExpression *ma, const FScriptPosition &pos)
 : FxExpression(pos)
 {
+	EmitTail = false;
 	if (mi != NULL && ma != NULL)
 	{
 		min = new FxIntCast(mi);
@@ -2679,6 +2744,18 @@ FxRandom::~FxRandom()
 //
 //==========================================================================
 
+PPrototype *FxRandom::ReturnProto()
+{
+	EmitTail = true;
+	return FxExpression::ReturnProto();
+}
+
+//==========================================================================
+//
+//
+//
+//==========================================================================
+
 FxExpression *FxRandom::Resolve(FCompileContext &ctx)
 {
 	CHECKRESOLVED();
@@ -2706,12 +2783,12 @@ int DecoRandom(VMFrameStack *stack, VMValue *param, int numparam, VMReturn *ret,
 	FRandom *rng = reinterpret_cast<FRandom *>(param[0].a);
 	if (numparam == 1)
 	{
-		ret->SetInt((*rng)());
+		ACTION_RETURN_INT((*rng)());
 	}
 	else if (numparam == 2)
 	{
 		int maskval = param[1].i;
-		ret->SetInt(rng->Random2(maskval));
+		ACTION_RETURN_INT(rng->Random2(maskval));
 	}
 	else if (numparam == 3)
 	{
@@ -2720,9 +2797,11 @@ int DecoRandom(VMFrameStack *stack, VMValue *param, int numparam, VMReturn *ret,
 		{
 			swapvalues(max, min);
 		}
-		ret->SetInt((*rng)(max - min + 1) + min);
+		ACTION_RETURN_INT((*rng)(max - min + 1) + min);
 	}
-	return 1;
+
+	// Shouldn't happen
+	return 0;
 }
 
 ExpEmit FxRandom::Emit(VMFunctionBuilder *build)
@@ -2735,17 +2814,27 @@ ExpEmit FxRandom::Emit(VMFunctionBuilder *build)
 	assert(((PSymbolVMFunction *)sym)->Function != NULL);
 	callfunc = ((PSymbolVMFunction *)sym)->Function;
 
+	int opcode = (EmitTail ? OP_TAIL_K : OP_CALL_K);
+
 	build->Emit(OP_PARAM, 0, REGT_POINTER | REGT_KONST, build->GetConstantAddress(rng, ATAG_RNG));
 	if (min != NULL && max != NULL)
 	{
 		EmitParameter(build, min, ScriptPosition);
 		EmitParameter(build, max, ScriptPosition);
-		build->Emit(OP_CALL_K, build->GetConstantAddress(callfunc, ATAG_OBJECT), 3, 1);
+		build->Emit(opcode, build->GetConstantAddress(callfunc, ATAG_OBJECT), 3, 1);
 	}
 	else
 	{
-		build->Emit(OP_CALL_K, build->GetConstantAddress(callfunc, ATAG_OBJECT), 1, 1);
+		build->Emit(opcode, build->GetConstantAddress(callfunc, ATAG_OBJECT), 1, 1);
 	}
+
+	if (EmitTail)
+	{
+		ExpEmit call;
+		call.Final = true;
+		return call;
+	}
+
 	ExpEmit out(build, REGT_INT);
 	build->Emit(OP_RESULT, 0, REGT_INT, out.RegNum);
 	return out;
@@ -2952,13 +3041,12 @@ int DecoFRandom(VMFrameStack *stack, VMValue *param, int numparam, VMReturn *ret
 		{
 			swapvalues(max, min);
 		}
-		ret->SetFloat(frandom * (max - min) + min);
+		ACTION_RETURN_FLOAT(frandom * (max - min) + min);
 	}
 	else
 	{
-		ret->SetFloat(frandom);
+		ACTION_RETURN_FLOAT(frandom);
 	}
-	return 1;
 }
 
 ExpEmit FxFRandom::Emit(VMFunctionBuilder *build)
@@ -2971,17 +3059,27 @@ ExpEmit FxFRandom::Emit(VMFunctionBuilder *build)
 	assert(((PSymbolVMFunction *)sym)->Function != NULL);
 	callfunc = ((PSymbolVMFunction *)sym)->Function;
 
+	int opcode = (EmitTail ? OP_TAIL_K : OP_CALL_K);
+
 	build->Emit(OP_PARAM, 0, REGT_POINTER | REGT_KONST, build->GetConstantAddress(rng, ATAG_RNG));
 	if (min != NULL && max != NULL)
 	{
 		EmitParameter(build, min, ScriptPosition);
 		EmitParameter(build, max, ScriptPosition);
-		build->Emit(OP_CALL_K, build->GetConstantAddress(callfunc, ATAG_OBJECT), 3, 1);
+		build->Emit(opcode, build->GetConstantAddress(callfunc, ATAG_OBJECT), 3, 1);
 	}
 	else
 	{
-		build->Emit(OP_CALL_K, build->GetConstantAddress(callfunc, ATAG_OBJECT), 1, 1);
+		build->Emit(opcode, build->GetConstantAddress(callfunc, ATAG_OBJECT), 1, 1);
 	}
+
+	if (EmitTail)
+	{
+		ExpEmit call;
+		call.Final = true;
+		return call;
+	}
+
 	ExpEmit out(build, REGT_FLOAT);
 	build->Emit(OP_RESULT, 0, REGT_FLOAT, out.RegNum);
 	return out;
@@ -2996,6 +3094,7 @@ ExpEmit FxFRandom::Emit(VMFunctionBuilder *build)
 FxRandom2::FxRandom2(FRandom *r, FxExpression *m, const FScriptPosition &pos)
 : FxExpression(pos)
 {
+	EmitTail = false;
 	rng = r;
 	if (m) mask = new FxIntCast(m);
 	else mask = new FxConstant(-1, pos);
@@ -3019,6 +3118,18 @@ FxRandom2::~FxRandom2()
 //
 //==========================================================================
 
+PPrototype *FxRandom2::ReturnProto()
+{
+	EmitTail = true;
+	return FxExpression::ReturnProto();
+}
+
+//==========================================================================
+//
+//
+//
+//==========================================================================
+
 FxExpression *FxRandom2::Resolve(FCompileContext &ctx)
 {
 	CHECKRESOLVED();
@@ -3042,9 +3153,19 @@ ExpEmit FxRandom2::Emit(VMFunctionBuilder *build)
 	assert(((PSymbolVMFunction *)sym)->Function != NULL);
 	callfunc = ((PSymbolVMFunction *)sym)->Function;
 
+	int opcode = (EmitTail ? OP_TAIL_K : OP_CALL_K);
+
 	build->Emit(OP_PARAM, 0, REGT_POINTER | REGT_KONST, build->GetConstantAddress(rng, ATAG_RNG));
 	EmitParameter(build, mask, ScriptPosition);
-	build->Emit(OP_CALL_K, build->GetConstantAddress(callfunc, ATAG_OBJECT), 2, 1);
+	build->Emit(opcode, build->GetConstantAddress(callfunc, ATAG_OBJECT), 2, 1);
+
+	if (EmitTail)
+	{
+		ExpEmit call;
+		call.Final = true;
+		return call;
+	}
+
 	ExpEmit out(build, REGT_INT);
 	build->Emit(OP_RESULT, 0, REGT_INT, out.RegNum);
 	return out;
@@ -3623,6 +3744,7 @@ FxActionSpecialCall::FxActionSpecialCall(FxExpression *self, int special, FArgum
 	Self = self;
 	Special = special;
 	ArgList = args;
+	EmitTail = false;
 }
 
 //==========================================================================
@@ -3643,6 +3765,18 @@ FxActionSpecialCall::~FxActionSpecialCall()
 //
 //==========================================================================
 
+PPrototype *FxActionSpecialCall::ReturnProto()
+{
+	EmitTail = true;
+	return FxExpression::ReturnProto();
+}
+
+//==========================================================================
+//
+//
+//
+//==========================================================================
+
 FxExpression *FxActionSpecialCall::Resolve(FCompileContext& ctx)
 {
 	CHECKRESOLVED();
@@ -3695,7 +3829,6 @@ FxExpression *FxActionSpecialCall::Resolve(FCompileContext& ctx)
 int DecoCallLineSpecial(VMFrameStack *stack, VMValue *param, int numparam, VMReturn *ret, int numret)
 {
 	assert(numparam > 2 && numparam < 8);
-	assert(numret == 1);
 	assert(param[0].Type == REGT_INT);
 	assert(param[1].Type == REGT_POINTER);
 	int v[5] = { 0 };
@@ -3704,8 +3837,7 @@ int DecoCallLineSpecial(VMFrameStack *stack, VMValue *param, int numparam, VMRet
 	{
 		v[i - 2] = param[i].i;
 	}
-	ret->SetInt(P_ExecuteSpecial(param[0].i, NULL, reinterpret_cast<AActor*>(param[1].a), false, v[0], v[1], v[2], v[3], v[4]));
+	ACTION_RETURN_INT(P_ExecuteSpecial(param[0].i, NULL, reinterpret_cast<AActor*>(param[1].a), false, v[0], v[1], v[2], v[3], v[4]));
-	return 1;
 }
 
 ExpEmit FxActionSpecialCall::Emit(VMFunctionBuilder *build)
@@ -3750,6 +3882,14 @@ ExpEmit FxActionSpecialCall::Emit(VMFunctionBuilder *build)
 	assert(((PSymbolVMFunction *)sym)->Function != NULL);
 	callfunc = ((PSymbolVMFunction *)sym)->Function;
 
+	if (EmitTail)
+	{
+		build->Emit(OP_TAIL_K, build->GetConstantAddress(callfunc, ATAG_OBJECT), 2 + i, 0);
+		ExpEmit call;
+		call.Final = true;
+		return call;
+	}
+
 	ExpEmit dest(build, REGT_INT);
 	build->Emit(OP_CALL_K, build->GetConstantAddress(callfunc, ATAG_OBJECT), 2 + i, 1);
 	build->Emit(OP_RESULT, 0, REGT_INT, dest.RegNum);
@@ -3767,6 +3907,7 @@ FxVMFunctionCall::FxVMFunctionCall(PFunction *func, FArgumentList *args, const F
 {
 	Function = func;
 	ArgList = args;
+	EmitTail = false;
 }
 
 //==========================================================================
@@ -3780,6 +3921,38 @@ FxVMFunctionCall::~FxVMFunctionCall()
 	SAFE_DELETE(ArgList);
 }
 
+//==========================================================================
+//
+//
+//
+//==========================================================================
+
+PPrototype *FxVMFunctionCall::ReturnProto()
+{
+	EmitTail = true;
+	return Function->Variants[0].Implementation->Proto;
+}
+
+//==========================================================================
+//
+//
+//
+//==========================================================================
+
+VMFunction *FxVMFunctionCall::GetDirectFunction()
+{
+	// If this return statement calls a function with no arguments,
+	// then it can be a "direct" function. That is, the DECORATE
+	// definition can call that function directly without wrapping
+	// it inside VM code.
+	if (EmitTail && (ArgList ? ArgList->Size() : 0) == 0 && (Function->Flags & VARF_Action))
+	{
+		return Function->Variants[0].Implementation;
+	}
+	
+	return nullptr;
+}
+
 //==========================================================================
 //
 // FxVMFunctionCall :: Resolve
@@ -3809,6 +3982,11 @@ FxExpression *FxVMFunctionCall::Resolve(FCompileContext& ctx)
 	{
 		ValueType = rets[0];
 	}
+	else
+	{
+		ValueType = TypeVoid;
+	}
+
 	return this;
 }
 
@@ -3821,19 +3999,14 @@ FxExpression *FxVMFunctionCall::Resolve(FCompileContext& ctx)
 //==========================================================================
 
 ExpEmit FxVMFunctionCall::Emit(VMFunctionBuilder *build)
-{
-	return Emit(build, false);
-}
-
-ExpEmit FxVMFunctionCall::Emit(VMFunctionBuilder *build, bool tailcall)
 {
 	assert(build->Registers[REGT_POINTER].GetMostUsed() >= 3);
-	int count = GetArgCount();
+	int count = (ArgList ? ArgList->Size() : 0);
 
 	if (count == 1)
 	{
 		ExpEmit reg;
-		if (CheckEmitCast(build, tailcall, reg))
+		if (CheckEmitCast(build, EmitTail, reg))
 		{
 			return reg;
 		}
@@ -3862,10 +4035,12 @@ ExpEmit FxVMFunctionCall::Emit(VMFunctionBuilder *build, bool tailcall)
 	VMFunction *vmfunc = Function->Variants[0].Implementation;
 	int funcaddr = build->GetConstantAddress(vmfunc, ATAG_OBJECT);
 	// Emit the call
-	if (tailcall)
+	if (EmitTail)
 	{ // Tail call
 		build->Emit(OP_TAIL_K, funcaddr, count, 0);
-		return ExpEmit();
+		ExpEmit call;
+		call.Final = true;
+		return call;
 	}
 	else if (vmfunc->Proto->ReturnTypes.Size() > 0)
 	{ // Call, expecting one result
@@ -3913,6 +4088,7 @@ bool FxVMFunctionCall::CheckEmitCast(VMFunctionBuilder *build, bool returnit, Ex
 				where.Free(build);
 			}
 			reg = ExpEmit();
+			reg.Final = true;
 		}
 		else
 		{
@@ -4540,54 +4716,77 @@ ExpEmit FxJumpStatement::Emit(VMFunctionBuilder *build)
 //
 //==========================================================================
 
-FxReturnStatement::FxReturnStatement(FxVMFunctionCall *call, const FScriptPosition &pos)
+FxReturnStatement::FxReturnStatement(FxExpression *value, const FScriptPosition &pos)
-: FxExpression(pos), Call(call)
+: FxExpression(pos), Value(value)
 {
+	ValueType = TypeVoid;
 }
 
 FxReturnStatement::~FxReturnStatement()
 {
-	SAFE_DELETE(Call);
+	SAFE_DELETE(Value);
 }
 
 FxExpression *FxReturnStatement::Resolve(FCompileContext &ctx)
 {
 	CHECKRESOLVED();
-	if (Call != NULL)
+	SAFE_RESOLVE_OPT(Value, ctx);
+
+	PPrototype *retproto;
+	if (Value == nullptr)
 	{
-		Call = static_cast<FxVMFunctionCall *>(Call->Resolve(ctx));
+		TArray<PType *> none(0);
-		ABORT(Call);
+		retproto = NewPrototype(none, none);
 	}
+	else
+	{
+		retproto = Value->ReturnProto();
+	}
+
+	ctx.CheckReturn(retproto, ScriptPosition);
+
 	return this;
 }
 
 ExpEmit FxReturnStatement::Emit(VMFunctionBuilder *build)
 {
-	// If we return nothing, use a regular RET opcode. If we return
+	// If we return nothing, use a regular RET opcode.
-	// something, use TAIL to call the function. Our return type
+	// Otherwise just return the value we're given.
-	// should be compatible with the called function's return type.
+	if (Value == nullptr)
-	if (Call == NULL)
 	{
 		build->Emit(OP_RET, RET_FINAL, REGT_NIL, 0);
 	}
 	else
 	{
-		Call->Emit(build, true);
+		ExpEmit ret = Value->Emit(build);
+
+		// Check if it is a function call that simplified itself
+		// into a tail call in which case we don't emit anything.
+		if (!ret.Final)
+		{
+			if (Value->ValueType == TypeVoid)
+			{ // Nothing is returned.
+				build->Emit(OP_RET, RET_FINAL, REGT_NIL, 0);
 			}
-	return ExpEmit();
+			else
+			{
+				build->Emit(OP_RET, RET_FINAL, ret.RegType | (ret.Konst ? REGT_KONST : 0), ret.RegNum);
+			}
+		}
+	}
+
+	ExpEmit out;
+	out.Final = true;
+	return out;
 }
 
 VMFunction *FxReturnStatement::GetDirectFunction()
 {
-	// If this return statement calls a function with no arguments,
+	if (Value != nullptr)
-	// then it can be a "direct" function. That is, the DECORATE
-	// definition can call that function directly without wrapping
-	// it inside VM code.
-	if (Call != NULL && Call->GetArgCount() == 0 && (Call->GetFunction()->Flags & VARF_Action))
 	{
-		return Call->GetVMFunction();
+		return Value->GetDirectFunction();
 	}
-	return NULL;
+	return nullptr;
 }
 
 //==========================================================================
@@ -4754,6 +4953,7 @@ FxExpression *FxStateByIndex::Resolve(FCompileContext &ctx)
 FxRuntimeStateIndex::FxRuntimeStateIndex(FxExpression *index)
 : FxExpression(index->ScriptPosition), Index(index)
 {
+	EmitTail = false;
 	ValueType = TypeState;
 }
 
@@ -4762,6 +4962,12 @@ FxRuntimeStateIndex::~FxRuntimeStateIndex()
 	SAFE_DELETE(Index);
 }
 
+PPrototype *FxRuntimeStateIndex::ReturnProto()
+{
+	EmitTail = true;
+	return FxExpression::ReturnProto();
+}
+
 FxExpression *FxRuntimeStateIndex::Resolve(FCompileContext &ctx)
 {
 	CHECKRESOLVED();
@@ -4794,14 +5000,12 @@ static int DecoHandleRuntimeState(VMFrameStack *stack, VMValue *param, int numpa
 		// Null is returned if the location was invalid which means that no jump will be performed
 		// if used as return value
 		// 0 always meant the same thing so we handle it here for compatibility
-		ret->SetPointer(nullptr, ATAG_STATE);
+		ACTION_RETURN_STATE(nullptr);
 	}
 	else
 	{
-		ret->SetPointer(stateinfo->mCallingState + index, ATAG_STATE);
+		ACTION_RETURN_STATE(stateinfo->mCallingState + index);
 	}
-
-	return 1;
 }
 
 ExpEmit FxRuntimeStateIndex::Emit(VMFunctionBuilder *build)
@@ -4823,8 +5027,16 @@ ExpEmit FxRuntimeStateIndex::Emit(VMFunctionBuilder *build)
 	assert(((PSymbolVMFunction *)sym)->Function != nullptr);
 	callfunc = ((PSymbolVMFunction *)sym)->Function;
 
+	if (EmitTail)
+	{
+		build->Emit(OP_TAIL_K, build->GetConstantAddress(callfunc, ATAG_OBJECT), 3, 1);
+		out.Final = true;
+	}
+	else
+	{
 		build->Emit(OP_CALL_K, build->GetConstantAddress(callfunc, ATAG_OBJECT), 3, 1);
 		build->Emit(OP_RESULT, 0, REGT_POINTER, out.RegNum);
+	}
 
 	return out;
 }

src/thingdef/thingdef_states.cpp

@@ -315,10 +315,10 @@ do_stop:
 				}
 
 				bool hasfinalret;
-				tcall->Code = ParseActions(sc, state, statestring, bag, tcall->Proto, hasfinalret);
+				tcall->Code = ParseActions(sc, state, statestring, bag, hasfinalret);
-				if (!hasfinalret)
+				if (!hasfinalret && tcall->Code != nullptr)
 				{
-					AddImplicitReturn(static_cast<FxSequence*>(tcall->Code), tcall->Proto, sc);
+					static_cast<FxSequence *>(tcall->Code)->Add(new FxReturnStatement(nullptr, sc));
 				}
 				goto endofstate;
 			}
@@ -351,125 +351,37 @@ endofstate:
 	sc.SetEscape(true);	// re-enable escape sequences
 }
 
-//==========================================================================
-//
-// AddImplicitReturn
-//
-// Adds an implied return; statement to the end of a code sequence.
-//
-//==========================================================================
-
-void AddImplicitReturn(FxSequence *code, const PPrototype *proto, FScanner &sc)
-{
-	if (code == NULL)
-	{
-		return;
-	}
-	if (proto == NULL || proto->ReturnTypes.Size() == 0)
-	{ // Returns nothing. Good. We can safely add an implied return.
-		code->Add(new FxReturnStatement(NULL, sc));
-	}
-	else
-	{ // Something was returned earlier in the sequence. Make it an error
-	  // instead of adding an implicit one.
-		sc.ScriptError("Not all paths return a value");
-	}
-}
-
-//==========================================================================
-//
-// ReturnCheck
-//
-// If proto1 is NULL, returns proto2. If proto2 is NULL, returns proto1.
-// If neither is null, checks if both prototypes define the same return
-// types. If not, an error is flagged.
-//
-//==========================================================================
-
-static PPrototype *ReturnCheck(PPrototype *proto1, PPrototype *proto2, FScanner &sc)
-{
-	if (proto1 == NULL)
-	{
-		return proto2;
-	}
-	if (proto2 == NULL)
-	{
-		return proto1;
-	}
-	// A prototype that defines fewer return types can be compatible with
-	// one that defines more if the shorter one matches the initial types
-	// for the longer one.
-	if (proto2->ReturnTypes.Size() < proto1->ReturnTypes.Size())
-	{ // Make proto1 the shorter one to avoid code duplication below.
-		swapvalues(proto1, proto2);
-	}
-	// If one prototype returns nothing, they both must.
-	if (proto1->ReturnTypes.Size() == 0)
-	{
-		if (proto2->ReturnTypes.Size() == 0)
-		{
-			return proto1;
-		}
-		proto1 = NULL;
-	}
-	else
-	{
-		for (unsigned i = 0; i < proto1->ReturnTypes.Size(); ++i)
-		{
-			if (proto1->ReturnTypes[i] != proto2->ReturnTypes[i])
-			{ // Incompatible
-				proto1 = NULL;
-				break;
-			}
-		}
-	}
-	if (proto1 == NULL)
-	{
-		sc.ScriptError("Return types are incompatible");
-	}
-	return proto1;
-}
-
 //==========================================================================
 //
 // ParseActions
 //
-// If this action block contains any return statements, the prototype for
-// one of them will be returned. This is used for deducing the return type
-// of anonymous functions. All called functions passed to return must have
-// matching return types.
-//
 //==========================================================================
 
-static FxExpression *ParseIf(FScanner &sc, FState state, FString statestring, Baggage &bag,
+static FxExpression *ParseIf(FScanner &sc, FState state, FString statestring, Baggage &bag, bool &lastwasret)
-	PPrototype *&retproto, bool &lastwasret)
 {
 	FxExpression *add, *cond;
-	FxExpression *true_part, *false_part = NULL;
+	FxExpression *true_part, *false_part = nullptr;
-	PPrototype *true_proto, *false_proto = NULL;
 	bool true_ret, false_ret = false;
 	sc.MustGetStringName("(");
 	cond = ParseExpression(sc, bag.Info);
 	sc.MustGetStringName(")");
 	sc.MustGetStringName("{");	// braces are mandatory
-	true_part = ParseActions(sc, state, statestring, bag, true_proto, true_ret);
+	true_part = ParseActions(sc, state, statestring, bag, true_ret);
 	sc.MustGetString();
 	if (sc.Compare("else"))
 	{
 		if (sc.CheckString("if"))
 		{
-			false_part = ParseIf(sc, state, statestring, bag, false_proto, false_ret);
+			false_part = ParseIf(sc, state, statestring, bag, false_ret);
 		}
 		else
 		{
 			sc.MustGetStringName("{");	// braces are still mandatory
-			false_part = ParseActions(sc, state, statestring, bag, false_proto, false_ret);
+			false_part = ParseActions(sc, state, statestring, bag, false_ret);
 			sc.MustGetString();
 		}
 	}
 	add = new FxIfStatement(cond, true_part, false_part, sc);
-	retproto = ReturnCheck(retproto, true_proto, sc);
-	retproto = ReturnCheck(retproto, false_proto, sc);
 	// If one side does not end with a return, we don't consider the if statement
 	// to end with a return. If the else case is missing, it can never be considered
 	// as ending with a return.
@@ -480,11 +392,9 @@ static FxExpression *ParseIf(FScanner &sc, FState state, FString statestring, Ba
 	return add;
 }
 
-static FxExpression *ParseWhile(FScanner &sc, FState state, FString statestring, Baggage &bag,
+static FxExpression *ParseWhile(FScanner &sc, FState state, FString statestring, Baggage &bag, bool &lastwasret)
-	PPrototype *&retproto, bool &lastwasret)
 {
 	FxExpression *cond, *code;
-	PPrototype *proto;
 	bool ret;
 
 	sc.MustGetStringName("(");
@@ -492,24 +402,21 @@ static FxExpression *ParseWhile(FScanner &sc, FState state, FString statestring,
 	sc.MustGetStringName(")");
 	sc.MustGetStringName("{"); // Enforce braces like for if statements.
 
-	code = ParseActions(sc, state, statestring, bag, proto, ret);
+	code = ParseActions(sc, state, statestring, bag, ret);
 	sc.MustGetString();
 
-	retproto = ReturnCheck(retproto, proto, sc);
 	lastwasret = false; // A while loop always jumps back.
 
 	return new FxWhileLoop(cond, code, sc);
 }
 
-static FxExpression *ParseDoWhile(FScanner &sc, FState state, FString statestring, Baggage &bag,
+static FxExpression *ParseDoWhile(FScanner &sc, FState state, FString statestring, Baggage &bag, bool &lastwasret)
-	PPrototype *&retproto, bool &lastwasret)
 {
 	FxExpression *cond, *code;
-	PPrototype *proto;
 	bool ret;
 
 	sc.MustGetStringName("{"); // Enforce braces like for if statements.
-	code = ParseActions(sc, state, statestring, bag, proto, ret);
+	code = ParseActions(sc, state, statestring, bag, ret);
 
 	sc.MustGetStringName("while");
 	sc.MustGetStringName("(");
@@ -518,20 +425,17 @@ static FxExpression *ParseDoWhile(FScanner &sc, FState state, FString statestrin
 	sc.MustGetStringName(";");
 	sc.MustGetString();
 
-	retproto = ReturnCheck(retproto, proto, sc);
 	lastwasret = false;
 
 	return new FxDoWhileLoop(cond, code, sc);
 }
 
-static FxExpression *ParseFor(FScanner &sc, FState state, FString statestring, Baggage &bag,
+static FxExpression *ParseFor(FScanner &sc, FState state, FString statestring, Baggage &bag, bool &lastwasret)
-	PPrototype *&retproto, bool &lastwasret)
 {
 	FxExpression *init = nullptr;
 	FxExpression *cond = nullptr;
 	FxExpression *iter = nullptr;
 	FxExpression *code = nullptr;
-	PPrototype *proto;
 	bool ret;
 
 	// Parse the statements.
@@ -560,24 +464,21 @@ static FxExpression *ParseFor(FScanner &sc, FState state, FString statestring, B
 
 	// Now parse the loop's content.
 	sc.MustGetStringName("{"); // Enforce braces like for if statements.
-	code = ParseActions(sc, state, statestring, bag, proto, ret);
+	code = ParseActions(sc, state, statestring, bag, ret);
 	sc.MustGetString();
 
-	retproto = ReturnCheck(retproto, proto, sc);
 	lastwasret = false;
 
 	return new FxForLoop(init, cond, iter, code, sc);
 }
 
-FxExpression *ParseActions(FScanner &sc, FState state, FString statestring, Baggage &bag,
+FxExpression *ParseActions(FScanner &sc, FState state, FString statestring, Baggage &bag, bool &endswithret)
-						   PPrototype *&retproto, bool &endswithret)
 {
 	// If it's not a '{', then it should be a single action.
 	// Otherwise, it's a sequence of actions.
 	if (!sc.Compare("{"))
 	{
 		FxVMFunctionCall *call = ParseAction(sc, state, statestring, bag);
-		retproto = call->GetVMFunction()->Proto;
 		endswithret = true;
 		return new FxReturnStatement(call, sc);
 	}
@@ -585,7 +486,6 @@ FxExpression *ParseActions(FScanner &sc, FState state, FString statestring, Bagg
 	const FScriptPosition pos(sc);
 
 	FxSequence *seq = NULL;
-	PPrototype *proto = NULL;
 	bool lastwasret = false;
 
 	sc.MustGetString();
@@ -595,38 +495,31 @@ FxExpression *ParseActions(FScanner &sc, FState state, FString statestring, Bagg
 		lastwasret = false;
 		if (sc.Compare("if"))
 		{ // Handle an if statement
-			add = ParseIf(sc, state, statestring, bag, proto, lastwasret);
+			add = ParseIf(sc, state, statestring, bag, lastwasret);
 		}
 		else if (sc.Compare("while"))
 		{ // Handle a while loop
-			add = ParseWhile(sc, state, statestring, bag, proto, lastwasret);
+			add = ParseWhile(sc, state, statestring, bag, lastwasret);
 		}
 		else if (sc.Compare("do"))
 		{ // Handle a do-while loop
-			add = ParseDoWhile(sc, state, statestring, bag, proto, lastwasret);
+			add = ParseDoWhile(sc, state, statestring, bag, lastwasret);
 		}
 		else if (sc.Compare("for"))
 		{ // Handle a for loop
-			add = ParseFor(sc, state, statestring, bag, proto, lastwasret);
+			add = ParseFor(sc, state, statestring, bag, lastwasret);
 		}
 		else if (sc.Compare("return"))
 		{ // Handle a return statement
 			lastwasret = true;
-			FxVMFunctionCall *retexp = NULL;
+			FxExpression *retexp = nullptr;
-			PPrototype *retproto;
 			sc.MustGetString();
 			if (!sc.Compare(";"))
 			{
-				retexp = ParseAction(sc, state, statestring, bag);
+				sc.UnGet();
+				retexp = ParseExpression(sc, bag.Info);
 				sc.MustGetStringName(";");
-				retproto = retexp->GetVMFunction()->Proto;
 			}
-			else
-			{ // Returning nothing; we still need a prototype for that.
-				TArray<PType *> notypes(0);
-				retproto = NewPrototype(notypes, notypes);
-			}
-			proto = ReturnCheck(proto, retproto, sc);
 			sc.MustGetString();
 			add = new FxReturnStatement(retexp, sc);
 		}
@@ -660,7 +553,6 @@ FxExpression *ParseActions(FScanner &sc, FState state, FString statestring, Bagg
 		}
 	}
 	endswithret = lastwasret;
-	retproto = proto;
 	return seq;
 }