1 files changed, 93 insertions, 22 deletions

diff --git a/src/cljcc/compiler.clj b/src/cljcc/compiler.clj
index 028cfbe..5d21d64 100644
--- a/src/cljcc/compiler.clj
+++ b/src/cljcc/compiler.clj
@@ -4,6 +4,10 @@
             [clojure.pprint :as pp]
             [cljcc.tacky :as t]))
 
+(def registers #{:ax :dx :r10 :r11})
+
+;;;; Instructions
+
 (defn- mov-instruction [src dst]
   {:op  :mov
    :src src
@@ -21,6 +25,21 @@
    :unary-operator unary-operator
    :operand operand})
 
+(defn- binary-instruction [binop src dst]
+  {:op :binary
+   :binary-operator binop
+   :src src
+   :dst dst})
+
+(defn- cdq-instruction []
+  {:op :cdq})
+
+(defn- idiv-instruction [operand]
+  {:op :idiv
+   :operand operand})
+
+;;;; Operands
+
 (defn- imm-operand [v]
   {:operand :imm
    :value v})
@@ -37,6 +56,8 @@
   {:operand :pseudo
    :identifier identifier})
 
+;;;; Tacky -> Instructions
+
 (defn- tacky-val->assembly-operand [val]
   (let [type (:type val)
         v (:value val)]
@@ -56,9 +77,28 @@
         unop (:unary-operator instruction)]
     [(mov-instruction src dst) (unary-instruction unop dst)]))
 
+(defn- tacky-binary->assembly [instruction]
+  (let [binop (:binary-operator instruction)
+        src1 (tacky-val->assembly-operand (:src1 instruction))
+        src2 (tacky-val->assembly-operand (:src2 instruction))
+        dst (tacky-val->assembly-operand (:dst instruction))
+        div? (= binop :div)
+        mod? (= binop :mod)]
+    (cond
+      div? [(mov-instruction src1 (reg-operand :ax))
+            (cdq-instruction)
+            (idiv-instruction src2)
+            (mov-instruction (reg-operand :ax) dst)]
+      mod? [(mov-instruction src1 (reg-operand :ax))
+            (cdq-instruction)
+            (idiv-instruction src2)
+            (mov-instruction (reg-operand :dx) dst)]
+      :else [(mov-instruction src1 dst) (binary-instruction binop src2 dst)])))
+
 (def tacky->assembly-transformers
   {:unary #'tacky-unary->assembly
-   :return #'tacky-return->assembly})
+   :return #'tacky-return->assembly
+   :binary #'tacky-binary->assembly})
 
 (defn- tacky-inst->assembly-inst [inst]
   (let [transformer-fn ((:type inst) tacky->assembly-transformers)]
@@ -112,30 +152,61 @@
     {:max-stack-val  (get pseudo-value-map "current")
      :instructions (map #(pseudo->stack-operand-instruction pseudo-value-map %) instructions)}))
 
-(defn- allocate-scratch-register [instruction]
-  (let [mov-both-stack? (fn [i]
-                          (and
-                           (= (:op i) :mov)
-                           (= (get-in i [:src :operand]) :stack)
-                           (= (get-in i [:dst :operand]) :stack)))
-        allocate-register-fn (fn [i]
-                               (if (mov-both-stack? i)
-                                 [(mov-instruction (get i :src) (reg-operand :r10))
-                                  (mov-instruction (reg-operand :r10) (get i :dst))]
-                                 i))]
-    (allocate-register-fn instruction)))
-
-(defn- fix-stack-instructions [{instructions :instructions max-stack-val :max-stack-val}]
-  (let [allocate-stack-inst (allocate-stack-instruction max-stack-val)
-        fixed-instructions (flatten (map allocate-scratch-register instructions))]
-    (cons allocate-stack-inst fixed-instructions)))
+(defn- fix-binary-instruction [instruction]
+  (let [binop (:binary-operator instruction)
+        src (:src instruction)
+        dst (:dst instruction)
+        mul? (= binop :mul)]
+    (if mul?
+      (let [dst-memory-address? (= :stack (:operand dst))]
+        (if dst-memory-address?
+          [(mov-instruction dst (reg-operand :r11))
+           (binary-instruction binop src (reg-operand :r11))
+           (mov-instruction (reg-operand :r11) dst)]
+          instruction))
+      (let [both-memory-address? (and
+                                  (= :stack (:operand src))
+                                  (= :stack (:operand dst)))]
+        (if both-memory-address?
+          [(mov-instruction src (reg-operand :r10))
+           (binary-instruction binop (reg-operand :r10) dst)]
+          instruction)))))
+
+(defn- fix-mov-instruction [instruction]
+  (let [both-memory-address? (and
+                              (= :stack (get-in instruction [:src :operand]))
+                              (= :stack (get-in instruction [:dst :operand])))]
+    (if both-memory-address?
+      [(mov-instruction (get instruction :src) (reg-operand :r10))
+       (mov-instruction (reg-operand :r10) (get instruction :dst))]
+      instruction)))
+
+(defn- fix-idiv-instruction [instruction]
+  (if (= :imm (get-in instruction [:operand :operand]))
+    [(mov-instruction (:operand instruction) (reg-operand :r10))
+     (idiv-instruction (reg-operand :r10))]
+    instruction))
+
+(def fix-instruction-map
+  {:idiv #'fix-idiv-instruction
+   :mov #'fix-mov-instruction
+   :binary #'fix-binary-instruction})
+
+(defn- fix-instruction [instruction]
+  (let [f (or ((:op instruction) fix-instruction-map) #'identity)]
+    (f instruction)))
+
+(defn- add-allocate-stack-instruction [{instructions :instructions max-stack-val :max-stack-val}]
+  (cons (allocate-stack-instruction max-stack-val) instructions))
 
 (defn- assembly-generate-instructions [tacky-instructions]
   (->> tacky-instructions
        (map tacky-inst->assembly-inst)
-       (flatten)
-       (replace-pseudoregisters)
-       (fix-stack-instructions)))
+       flatten
+       replace-pseudoregisters
+       add-allocate-stack-instruction
+       (map fix-instruction)
+       flatten))
 
 (defn- transform-function [_return-type identifier args body]
   {:op :function
@@ -165,7 +236,7 @@
   (pp/pprint
    (generate-assembly
     "int main(void) {
-return ~(-(~(-1)));
+return 1 + 2;
 }"))
 
   ())