From dca5d7dc70964276a01db3ac9d0fde00d7e7fc95 Mon Sep 17 00:00:00 2001
From: Hazem Krimi <krimihazem1@gmail.com>
Date: Mon, 13 May 2024 20:01:19 +0100
Subject: [PATCH] Fix return translation

---
 include/functions.h | 20 ++++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/include/functions.h b/include/functions.h
index 94afaa0..eeffabd 100644
--- a/include/functions.h
+++ b/include/functions.h
@@ -75,17 +75,24 @@ string translateCall(string name, int args) {
 
 string translateReturn() {
     stringstream output;
+    string cleanupLabel = "cleanup$ret" + generateRandomLabel(3);
+    string endLabel = "end$ret" + generateRandomLabel(3);
 
     output << "@LCL" << endl;
-    output << "D=A" << endl;
+    output << "D=M" << endl;
     output << "@END_FRAME" << endl;
     output << "M=D" << endl;
+
     output << "@RETURN_ADDR" << endl;
     output << "M=D" << endl;
     output << "@5" << endl;
     output << "D=A" << endl;
     output << "@RETURN_ADDR" << endl;
     output << "M=M-D" << endl;
+    output << "A=M" << endl;
+    output << "D=M" << endl;
+    output << "@RETURN_ADDR" << endl;
+    output << "M=D" << endl;
 
     output << "@ARG" << endl;
     output << "D=M" << endl;
@@ -108,6 +115,8 @@ string translateReturn() {
     output << "D=M" << endl;
     output << "@1" << endl;
     output << "D=D-A" << endl;
+    output << "A=D" << endl;
+    output << "D=M" << endl;
     output << "@THAT" << endl;
     output << "M=D" << endl;
     
@@ -115,6 +124,8 @@ string translateReturn() {
     output << "D=M" << endl;
     output << "@2" << endl;
     output << "D=D-A" << endl;
+    output << "A=D" << endl;
+    output << "D=M" << endl;
     output << "@THIS" << endl;
     output << "M=D" << endl;
     
@@ -122,6 +133,8 @@ string translateReturn() {
     output << "D=M" << endl;
     output << "@3" << endl;
     output << "D=D-A" << endl;
+    output << "A=D" << endl;
+    output << "D=M" << endl;
     output << "@ARG" << endl;
     output << "M=D" << endl;
 
@@ -129,11 +142,14 @@ string translateReturn() {
     output << "D=M" << endl;
     output << "@4" << endl;
     output << "D=D-A" << endl;
+    output << "A=D" << endl;
+    output << "D=M" << endl;
     output << "@LCL" << endl;
     output << "M=D" << endl;
 
     output << "@RETURN_ADDR" << endl;
-    output << "0;JMP" << endl;
+    output << "D=M" << endl;
+    output << "D;JMP" << endl;
 
     return output.str();
 }
\ No newline at end of file