From 1568a0caccc2c264a322dadf81bab16b68195f87 Mon Sep 17 00:00:00 2001
From: Frederick Yin <fkfd@fkfd.me>
Date: Tue, 23 Aug 2022 23:40:41 +0800
Subject: hack-vm: fix assembly for function commands

---
 projects/hack-vm/function.py | 41 +++++++++++++++++++++++------------------
 1 file changed, 23 insertions(+), 18 deletions(-)

(limited to 'projects/hack-vm')

diff --git a/projects/hack-vm/function.py b/projects/hack-vm/function.py
index 596c9cf..61c86f6 100644
--- a/projects/hack-vm/function.py
+++ b/projects/hack-vm/function.py
@@ -4,27 +4,27 @@ D=A
 A=M
 M=D
 @LCL
-D=A
+D=M
 @SP
 AM=M+1
 M=D
 @ARG
-D=A
+D=M
 @SP
 AM=M+1
 M=D
 @THIS
-D=A
+D=M
 @SP
 AM=M+1
 M=D
 @THAT
-D=A
+D=M
 @SP
 AM=M+1
 M=D
 @SP
-DM=M+1
+MD=M+1
 @LCL
 M=D
 @{arg_lcl_offset}
@@ -36,22 +36,26 @@ M=D
 ({return_address})
 """
 
-FUNCTION_ASM = """({function_label})
-@SP
-A=M
-M=0
-{init_vars}@SP
-M=M+1
-"""
+FUNCTION_ASM = "({function_label})\n"
 
-# repeated for each local variable of function
-INIT_VAR_ASM = """@SP
+# repeated (n-1) times for function with n arguments
+INIT_ONE_VAR_ASM = """@SP
 AM=M+1
 M=0
 """
 
-RETURN_ASM = """@ARG
+# concatenated only when function has arguments
+INIT_VARS_ASM = """@SP
 A=M
+M=0
+{init_vars}@SP
+M=M+1
+"""
+
+RETURN_ASM = """@LCL
+D=M
+@5
+A=D-A
 D=M
 @{tmp}
 M=D
@@ -112,9 +116,10 @@ def translate_function(action, function, n, prog, verbose=False):
             function_label=f"{function}",
         )
     else:
-        asm += FUNCTION_ASM.format(
-            function_label=f"{function}", init_vars=INIT_VAR_ASM * (n - 1)
-        )
+        asm += FUNCTION_ASM.format(function_label=f"{function}")
+        if n > 0:
+            # function has arguments
+            asm += INIT_VARS_ASM.format(init_vars=INIT_ONE_VAR_ASM * (n - 1))
 
     return asm
 
-- 
cgit v1.2.3