From 79ca8e78b315ca148ebadf4508748ede1fdf6969 Mon Sep 17 00:00:00 2001
From: Bananymous <bananymousosq@gmail.com>
Date: Mon, 29 Apr 2024 18:18:24 +0300
Subject: [PATCH] Optimize binary_op generation to use less stack if possible

---
 04_semantics_and_running/main.py | 67 +++++++++++++++++++++++++-------
 1 file changed, 54 insertions(+), 13 deletions(-)

diff --git a/04_semantics_and_running/main.py b/04_semantics_and_running/main.py
index c0048d7..f2851e3 100644
--- a/04_semantics_and_running/main.py
+++ b/04_semantics_and_running/main.py
@@ -487,23 +487,64 @@ def compile_ast(node: ASTnode, compile_data: CompileData) -> None:
             else:
                 assert node.child_rhs.type == 'int'
 
-            # calculate LHS and store it on stack
+
+            old_code = compile_data.code
+
+            # compile LHS
+            compile_data.code = []
             compile_ast(node.child_lhs, compile_data)
-            compile_data.code.append(Instruction('subq', ['$16', '%rsp']))
-            compile_data.code.append(Instruction('movq', ['%rax', '0(%rsp)']))
+            lhs_code = compile_data.code
 
-            # calculate RHS and store it in RCX
+            # compile RHS
+            compile_data.code = []
             compile_ast(node.child_rhs, compile_data)
-            if node.child_lhs.type == 'date' and node.child_rhs.type == 'int':
-                # multiply RHS by number of seconds in a day so we can perform arithmetic on dates
-                compile_data.code.append(Instruction('imulq', ['$86400', '%rax', '%rcx']))
-            else:
-                compile_data.code.append(Instruction('movq', ['%rax', '%rcx']))
+            rhs_code = compile_data.code
 
-            # prepare registers, RAX contains LHS and RCX contains RHS
-            # and restore restore stack
-            compile_data.code.append(Instruction('movq', ['0(%rsp)', '%rax']))
-            compile_data.code.append(Instruction('addq', ['$16', '%rsp']))
+            compile_data.code = old_code
+
+            # check if we can use temporary registers instead of stack
+            temp_registers = ['%rcx', '%r8', '%r9', '%r10', '%r11']
+            temp_reg = None
+            for reg in temp_registers:
+                valid = True
+                for instruction in lhs_code:
+                    if reg in instruction.operands:
+                        valid = False
+                        break
+                if valid:
+                    temp_reg = reg
+                    break
+
+            # check if lhs uses call, this determines whether we need to align stack
+            lhs_call = False
+            for instruction in lhs_code:
+                if instruction.opcode == 'call':
+                    lhs_call = True
+                    break
+
+            # Add code for RHS calculation
+            compile_data.code += rhs_code
+            if temp_reg is not None:
+                compile_data.code.append(Instruction('movq', ['%rax', temp_reg]))
+            elif not lhs_call:
+                compile_data.code.append(Instruction('pushq', ['%rax']))
+            else:
+                compile_data.code.append(Instruction('subq', ['$16', '%rsp']))
+                compile_data.code.append(Instruction('movq', ['%rax', '0(%rsp)']))
+
+            # Add code for LHS calculation
+            compile_data.code += lhs_code
+            if temp_reg is not None:
+                compile_data.code.append(Instruction('movq', [temp_reg, '%rcx']))
+            elif not lhs_call:
+                compile_data.code.append(Instruction('popq', ['%rcx']))
+            else:
+                compile_data.code.append(Instruction('movq', ['0(%rsp)', '%rcx']))
+                compile_data.code.append(Instruction('addq', ['$16', '%rsp']))
+
+            # If we are adding or subtracting dates with integers, multiply the integer by number of seconds in a day
+            if node.child_lhs.type == 'date' and node.child_rhs.type == 'int':
+                compile_data.code.append(Instruction('imulq', ['$86400', '%rcx']))
 
             # perform operation
             if node.value == '+':