hack-vm: fix assembly for function commands

1 files changed, 23 insertions, 18 deletions

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