From 221fc2361404ed0a668c7946ba771bdd361de754 Mon Sep 17 00:00:00 2001
From: Cassie Jones <code@witchoflight.com>
Date: Sun, 1 Mar 2020 17:28:31 +0100
Subject: [PATCH] Reduce stack traffic by finding loads through SSA

Instead of loading Load instructions directly and then storing them onto
the stack, we perform loads by looking through the SSA and seeing what
instruction generated them. If it was a load, we run that load directly,
otherwise we load from the stack slot that it corresponds to.
---
 .gitignore       |  1 +
 src/trans/x64.rs | 41 +++++++++++++++++++++--------------------
 2 files changed, 22 insertions(+), 20 deletions(-)

diff --git a/.gitignore b/.gitignore
index ea8c4bf..5aa569a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1 +1,2 @@
 /target
+a.out
diff --git a/src/trans/x64.rs b/src/trans/x64.rs
index f1952bc..6d3c921 100644
--- a/src/trans/x64.rs
+++ b/src/trans/x64.rs
@@ -201,26 +201,24 @@ fn compile_func(func: &Func) -> Definition {
     body.push(Inst::Push(Reg::Rbx));
     body.push(Inst::Mov(Reg::Rbx, Reg::Rdi));
     body.push(Inst::Sub(Reg::Rsp, stack_slots as u32 * 8));
+    let address = |v| match v {
+        Var::Global(x) => Addr::Rip(Label::Global(x)),
+        Var::Param(x) => param(x),
+        Var::Upvar(x) => upvar(x),
+        Var::Ssa(s) => stack(&s),
+    };
+    let load = |s| {
+        match func.block[s] {
+            Code::Load(v) => address(v),
+            _ => stack(s),
+        }
+    };
     for ssa in &func.order {
         match &func.block[ssa] {
-            Code::Load(Var::Global(x)) => {
-                body.push(Inst::Load(Reg::Rax, Addr::Rip(Label::Global(*x))));
-                body.push(Inst::Store(stack(ssa), Reg::Rax));
-            }
-            Code::Load(Var::Param(x)) => {
-                body.push(Inst::Load(Reg::Rax, param(*x)));
-                body.push(Inst::Store(stack(ssa), Reg::Rax));
-            }
-            Code::Load(Var::Upvar(x)) => {
-                body.push(Inst::Load(Reg::Rax, upvar(*x)));
-                body.push(Inst::Store(stack(ssa), Reg::Rax));
-            }
-            Code::Load(Var::Ssa(ssa2)) => {
-                body.push(Inst::Load(Reg::Rax, stack(ssa2)));
-                body.push(Inst::Store(stack(ssa), Reg::Rax));
-            }
+            // We don't need to generate code for loads themselves, they're consulted by other accesses
+            Code::Load(_) => (),
             Code::StoreGlobal(x, s) => {
-                body.push(Inst::Load(Reg::Rax, stack(s)));
+                body.push(Inst::Load(Reg::Rax, load(s)));
                 body.push(Inst::Store(Addr::Rip(Label::Global(*x)), Reg::Rax));
             }
             Code::MakeLam {
@@ -240,16 +238,16 @@ fn compile_func(func: &Func) -> Definition {
                 let mut terms = terms.iter();
                 body.push(Inst::Load(
                     Reg::Rdi,
-                    stack(terms.next().expect("a function to apply")),
+                    load(terms.next().expect("a function to apply")),
                 ));
                 match terms.next() {
-                    Some(s) => body.push(Inst::Load(Reg::Rsi, stack(s))),
+                    Some(s) => body.push(Inst::Load(Reg::Rsi, load(s))),
                     None => body.push(Inst::Imm(Reg::Rsi, 0)),
                 }
                 body.push(Inst::Call(extern_label!("ivy_app")));
                 while let Some(term) = terms.next() {
                     body.push(Inst::Mov(Reg::Rdi, Reg::Rax));
-                    body.push(Inst::Load(Reg::Rsi, stack(term)));
+                    body.push(Inst::Load(Reg::Rsi, load(term)));
                     body.push(Inst::Call(extern_label!("ivy_app_mut")));
                 }
                 body.push(Inst::Store(stack(ssa), Reg::Rax));
@@ -260,6 +258,9 @@ fn compile_func(func: &Func) -> Definition {
             }
         }
     }
+    if let Code::Load(v) = func.block[&func.result] {
+        body.push(Inst::Load(Reg::Rax, address(v)));
+    }
     body.push(Inst::Add(Reg::Rsp, stack_slots as u32 * 8));
     body.push(Inst::Pop(Reg::Rbx));
     body.push(Inst::Ret);
-- 
2.47.0