Refactor calling conventions to support correct representation of structures

2 files changed, 49 insertions, 12 deletions

diff --git a/lang/rust/pseudorust.cpp b/lang/rust/pseudorust.cpp
index 92e9e56a..2e233550 100644
--- a/lang/rust/pseudorust.cpp
+++ b/lang/rust/pseudorust.cpp
@@ -1900,6 +1900,24 @@ void PseudoRustFunction::GetExprText(const HighLevelILInstruction& instr, HighLe
 		}();
 		break;
 
+	case HLIL_PASS_BY_REF:
+		[&]() {
+			const auto srcExpr = instr.GetSourceExpr<HLIL_PASS_BY_REF>();
+			GetExprText(srcExpr, tokens, settings, UnaryOperatorPrecedence);
+			if (exprType != InnerExpression)
+				tokens.AppendSemicolon();
+		}();
+		break;
+
+	case HLIL_RETURN_BY_REF:
+		[&]() {
+			const auto srcExpr = instr.GetSourceExpr<HLIL_RETURN_BY_REF>();
+			GetExprText(srcExpr, tokens, settings, UnaryOperatorPrecedence);
+			if (exprType != InnerExpression)
+				tokens.AppendSemicolon();
+		}();
+		break;
+
 	case HLIL_FCMP_E:
 	case HLIL_CMP_E:
 		[&]() {
diff --git a/lang/rust/rusttypes.cpp b/lang/rust/rusttypes.cpp
index 70bf3162..0d966c3a 100644
--- a/lang/rust/rusttypes.cpp
+++ b/lang/rust/rusttypes.cpp
@@ -147,12 +147,13 @@ vector<InstructionTextToken> RustTypePrinter::GetTypeTokensAfterNameInternal(
 			vector<InstructionTextToken> paramTokens = GetTypeTokensAfterName(params[i].type.GetValue(), platform,
 				params[i].type.GetCombinedConfidence(baseConfidence), type, escaping);
 
-			if (functionHeader)
+			auto var = params[i].location.GetVariableForParameter(i);
+			if (functionHeader && var.has_value())
 			{
 				for (auto& token : paramTokens)
 				{
 					token.context = LocalVariableTokenContext;
-					token.address = params[i].location.ToIdentifier();
+					token.address = var->ToIdentifier();
 				}
 			}
 
@@ -164,40 +165,51 @@ vector<InstructionTextToken> RustTypePrinter::GetTypeTokensAfterNameInternal(
 			else
 			{
 				nameToken = InstructionTextToken(ArgumentNameToken, NameList::EscapeTypeName(params[i].name, escaping), i);
-				if (functionHeader)
+				if (functionHeader && var.has_value())
 				{
 					nameToken.context = LocalVariableTokenContext;
-					nameToken.address = params[i].location.ToIdentifier();
+					nameToken.address = var->ToIdentifier();
 				}
 			}
 			tokens.push_back(nameToken);
 			tokens.emplace_back(TextToken, ": ");
+
+			if (params[i].locationSource == PassByReferenceLocationSource || params[i].location.indirect)
+				tokens.emplace_back(baseConfidence, OperationToken, "&");
+
 			tokens.insert(tokens.end(), paramTokens.begin(), paramTokens.end());
 
-			if (!params[i].defaultLocation && platform)
+			if (params[i].locationSource == CustomLocationSource && platform && var.has_value())
 			{
-				switch (params[i].location.type)
+				switch (var->type)
 				{
 				case RegisterVariableSourceType:
 				{
-					string registerName = platform->GetArchitecture()->GetRegisterName((uint32_t)params[i].location.storage);
+					string registerName = platform->GetArchitecture()->GetRegisterName((uint32_t)var->storage);
 					tokens.emplace_back(TextToken, " @ ");
-					tokens.emplace_back(RegisterToken, NameList::EscapeTypeName(registerName, escaping), params[i].location.storage);
+					tokens.emplace_back(RegisterToken, NameList::EscapeTypeName(registerName, escaping), var->storage);
 					break;
 				}
 				case FlagVariableSourceType:
 				{
-					string flagName = platform->GetArchitecture()->GetFlagName((uint32_t)params[i].location.storage);
+					string flagName = platform->GetArchitecture()->GetFlagName((uint32_t)var->storage);
 					tokens.emplace_back(TextToken, " @ ");
-					tokens.emplace_back(AnnotationToken, NameList::EscapeTypeName(flagName, escaping), params[i].location.storage);
+					tokens.emplace_back(AnnotationToken, NameList::EscapeTypeName(flagName, escaping), var->storage);
 					break;
 				}
 				case StackVariableSourceType:
 				{
 					tokens.emplace_back(TextToken, " @ ");
 					char storageStr[32];
-					snprintf(storageStr, sizeof(storageStr), "%" PRIi64, params[i].location.storage);
-					tokens.emplace_back(IntegerToken, storageStr, params[i].location.storage);
+					snprintf(storageStr, sizeof(storageStr), "%" PRIi64, var->storage);
+					tokens.emplace_back(IntegerToken, storageStr, var->storage);
+					break;
+				}
+				default:
+				{
+					string locationStr = params[i].location.ToString(platform->GetArchitecture());
+					tokens.emplace_back(TextToken, " @ ");
+					tokens.emplace_back(ValueLocationToken, locationStr);
 					break;
 				}
 				}
@@ -230,6 +242,13 @@ vector<InstructionTextToken> RustTypePrinter::GetTypeTokensAfterNameInternal(
 					i.context = FunctionReturnTokenContext;
 			}
 			tokens.insert(tokens.end(), retn.begin(), retn.end());
+			if (!type->GetReturnValue().defaultLocation)
+			{
+				auto location = type->GetReturnValue().location;
+				string locationStr = location->ToString(platform->GetArchitecture());
+				tokens.emplace_back(TextToken, " @ ");
+				tokens.emplace_back(location.GetCombinedConfidence(baseConfidence), ValueLocationToken, locationStr);
+			}
 		}
 		break;
 	}