5303b4975854be55bfc5d36aedcf5e607ebe8db9
[scheme.git] / codegen.scm
1 (load "typecheck.scm")
2 (load "ast.scm")
3
4 (define (emit . s)
5   (begin
6     (apply printf s)
7     (display "\n")))
8
9 (define (codegen-add xs si env)
10   (define (go ys)
11     (if (null? ys)
12       (emit "movq ~a(%rbp), %rax" si)
13       (begin
14         (let ((y (car ys)))
15              (if (integer? y)
16                  (emit "addq $~a, ~a(%rbp)" y si)
17                  (begin
18                    (codegen-expr y (- si wordsize) env)
19                    (emit "addq %rax, ~a(%rbp)" si))))
20         (go (cdr ys)))))
21   (begin
22                                         ; use si(%rbp) as the accumulator
23     (emit "movq $0, ~a(%rbp)" si)
24     (go xs)))
25
26 (define (codegen-binop opcode)
27   (lambda (a b si env)
28     (codegen-expr b si env)
29     (emit "movq %rax, ~a(%rbp)" si)
30     (codegen-expr a (- si wordsize) env)
31     (emit "~a ~a(%rbp), %rax" opcode si)))
32
33 (define codegen-sub (codegen-binop "sub"))
34
35 (define codegen-mul (codegen-binop "imul"))
36
37 (define (codegen-not x si env)
38   (codegen-expr x si env)
39   (emit "notq %rax")
40   (emit "andq $1, %rax"))
41
42 (define (codegen-eq a b si env)
43   (codegen-expr a si env)
44   (emit "movq %rax, ~a(%rbp)" si)
45   (codegen-expr b (- si wordsize) env)
46   (emit "subq ~a(%rbp), %rax" si)
47   (emit "not %rax")
48   (emit "andq $1, %rax"))
49
50 (define (codegen-print x si env)
51   (codegen-expr x si env) ; x should be a static-string, producing a label
52
53                                         ; get the length of the null terminated string
54   (emit "mov %rax, %rdi")
55   (emit "xor %al, %al")   ; set %al to 0
56   (emit "mov $-1, %rcx") ; max search length = max int = -1
57   (emit "cld")           ; clear direction flag, search up in memory
58   (emit "repne scasb")   ; scan string, %rcx = -strlen - 1 - 1
59   
60   (emit "not %rcx")      ; -%rcx = strlen + 1
61   (emit "dec %rcx")
62   
63   (emit "mov %rcx, %rdx") ; number of bytes
64   (emit "mov %rax, %rsi") ; addr of string
65   (emit "mov $1, %rax") ; file handle 1 (stdout)
66   (emit "mov $1, %rdi") ; syscall 1 (write)
67   (emit "syscall"))
68
69 (define (range s n)
70   (if (= 0 n) '()
71       (append (range s (- n 1))
72               (list (+ s (- n 1))))))
73
74 (define wordsize 8)
75
76 (define (codegen-let bindings body si env)
77   (let* ((stack-offsets (map (lambda (x) (- si (* x wordsize)))
78                              (range 0 (length bindings))))
79          (inner-si (- si (* (length bindings) wordsize)))
80          (names (map car bindings))
81          (exprs (map cadr bindings))
82
83          ; recursive let bindings: build environment as we go
84          (inner-env (fold-left
85                      (lambda (env name expr offset)
86                        (codegen-expr expr inner-si env)
87                        (emit "movq %rax, ~a(%rbp)" offset)
88                        (cons (cons name offset) env))
89                      env names exprs stack-offsets)))
90     (for-each (lambda (form)
91                 (codegen-expr form inner-si inner-env))
92               body)))
93
94 (define (codegen-var name si env)
95   (when (not (assoc name env))
96     (error #f (format "Variable ~a is not bound" name)))
97   (let ((offset (cdr (assoc name env))))
98     (emit "movq ~a(%rbp), %rax" offset)))
99
100 (define cur-lambda 0)
101 (define (fresh-lambda)
102   (set! cur-lambda (+ 1 cur-lambda))
103   (format "_lambda~a" (- cur-lambda 1)))
104
105                                         ; for now we can only call closures
106 (define (codegen-call closure args si env)
107   (when (not (eq? (ast-type closure) 'closure))
108     (error #f (format "~a is not a closure" closure)))
109   (let* ((captured (caddr closure))
110          (label (cadr closure))
111          (argument-start (length captured)))
112
113                                         ; first move the captured variables into param registers
114     (for-each
115      (lambda (e i)
116        (emit "movq ~a(%rbp), ~a"
117              (cdr (assoc e env)) ; offset of the var
118              (param-register i)))
119      captured (range 0 (length captured)))
120
121
122                                         ; then codegen the arguments and move them into the next param registers
123     (for-each
124      (lambda (e i)
125        (begin
126          (codegen-expr e si env)
127                                         ; move result to correct param register
128          (emit "movq %rax, ~a" (param-register i))))
129      args (range argument-start (length args)))
130
131     (emit "addq $~a, %rsp" si) ; adjust the stack pointer to account all the stuff we put in the env
132     (emit "callq ~a" label)
133     (emit "subq $~a, %rsp" si)))
134
135 (define (codegen-lambda l)
136   (let* ((label (car l))
137          (args (cadr l))
138          (captured (caddr l))
139          (body (cadddr l))
140 ; params = what actually gets passed
141          (params (append captured args))
142
143          (param-registers (map param-register
144                                (range 0 (length params))))
145          (stack-offsets (map (lambda (i)
146                                (* (- wordsize) i))
147                              (range 1 (length params))))
148
149          (copy-insts (map (lambda (r o)
150                             (format "movq ~a, ~a(%rbp)" r o))
151                           param-registers stack-offsets))
152
153          (env (map cons params stack-offsets)))
154     (emit "~a:" label)
155     (display "## lambda body: ")
156     (display body)
157     (newline)
158     (display "## environment: ")
159     (display env)
160     (newline)
161
162     (emit "push %rbp") ; preserve caller's base pointer
163     (emit "movq %rsp, %rbp") ; set up our own base pointer
164     
165     (for-each emit copy-insts)
166     (codegen-expr body (* (- wordsize) (+ 1 (length params))) env)
167
168     (emit "pop %rbp") ; restore caller's base pointer
169     (emit "ret")))
170
171 (define cur-label 0)
172 (define (fresh-label)
173   (set! cur-label (+ 1 cur-label))
174   (format "label~a" (- cur-label 1)))
175
176 (define (codegen-if cond then else si env)
177   (codegen-expr cond si env)
178   (emit "cmpq $0, %rax")
179   (let ((exit-label (fresh-label))
180         (else-label (fresh-label)))
181     (emit "je ~a" else-label)
182     (codegen-expr then si env)
183     (emit "jmp ~a" exit-label)
184     (emit "~a:" else-label)
185     (codegen-expr else si env)
186     (emit "~a:" exit-label)))
187
188 (define (codegen-expr e si env)
189   (case (ast-type e)
190     ('builtin e)
191     ('closure e)
192     ('app
193      (let ((callee (codegen-expr (car e) si env)))
194        (case callee
195          ('+ (codegen-add (cdr e) si env))
196          ('- (codegen-sub (cadr e) (caddr e) si env))
197          ('* (codegen-mul (cadr e) (caddr e) si env))
198          ('! (codegen-not (cadr e) si env))
199          ('= (codegen-eq  (cadr e) (caddr e) si env))
200          ('bool->int (codegen-expr (cadr e) si env))
201          ('print (codegen-print (cadr e) si env))
202          (else (codegen-call callee (cdr e) si env)))))
203
204     ('let (codegen-let (let-bindings e)
205                        (let-body e)
206                        si
207                        env))
208
209     ('var (codegen-var e si env))
210
211     ('if (codegen-if (cadr e) (caddr e) (cadddr e) si env))
212     
213     ('bool-literal (emit "movq $~a, %rax" (if e 1 0)))
214     ('int-literal (emit "movq $~a, %rax" e))
215     
216     ('static-string (emit "movq $~a, %rax" (cadr e))) ; move label
217
218     (else (error #f "don't know how to codegen this"))))
219
220
221 (define (fold-map f x) (fold-left append '() (map f x)))
222
223 (define (free-vars prog)
224   (define bound '())
225   (define (collect e)
226     (case (ast-type e)
227       ('builtin '()) ; do nothing
228       ('var (if (memq e bound) '() (list e)))
229       ('lambda
230           (set! bound (append (lambda-args e) bound))
231         (collect (lambda-body e)))
232
233       ('app (fold-map collect e))
234       ('let
235           (let ((bind-fvs (fold-map (lambda (a)
236                                       ((set! bound (cons (car a) bound))
237                                        (collect (cdr a))))
238                                     (let-bindings cadr)))
239                 (body-fvs (fold-map collect (let-body e))))
240             (append bind-fvs body-fvs)))
241       (else '())))
242   (collect prog))
243
244                                         ; ((lambda (x) (+ x 1)) 42) => {lambda0: (x) (+ x 1)}, (@lambda0 42)
245 (define (extract-lambdas program)
246   (define lambdas '())
247   (define (add-lambda e)
248     (let* ((label (fresh-lambda))
249            (args (lambda-args e))
250            (captured (free-vars e))
251            (body (extract (lambda-body e)))
252            (new-lambda (list label args captured body)))
253       (set! lambdas (cons new-lambda lambdas))
254       `(closure ,label ,captured))) ; todo: should we string->symbol?
255   (define (extract e)
256     (case (ast-type e)
257       ('lambda (add-lambda e))
258       ('let `(let ,(map extract (let-bindings e))
259                ,@(map extract (let-body e))))
260       ('app (append (list (extract (car e)))
261                     (map extract (cdr e))))
262       (else (ast-traverse extract e))))
263   (let ((transformed (extract program)))
264     (cons lambdas transformed)))
265
266 (define (extract-strings program)
267   (let ((cur-string 0)
268         (strings '())) ; assoc list of labels -> string
269     (define (fresh-string)
270       (set! cur-string (+ cur-string 1))
271       (format "string~a" (- cur-string 1)))
272     (define (extract e)
273       (case (ast-type e)
274         ('string-literal
275          (let ((label (fresh-string)))
276            (set! strings (cons (cons label e) strings))
277            `(static-string ,label)))
278         (else (ast-traverse extract e))))
279     (let ((transformed (extract program)))
280       (cons strings transformed))))
281
282 (define (codegen-string-data s)
283   (emit "~a:" (car s))
284   (emit "\t.string \"~a\"" (cdr s)))
285
286 ;; (define (amd64-abi f)
287 ;;                                      ; preserve registers
288 ;;   (emit "push %rbp")
289 ;;   ;; (emit "push %rbx")
290 ;;   ;; (for-each (lambda (i)
291 ;;   ;;               (emit (string-append
292 ;;   ;;                      "push %r"
293 ;;   ;;                      (number->string i))))
294 ;;   ;;             '(12 13 14 15))
295
296 ;;   (emit "movq %rsp, %rbp")              ; set up the base pointer
297
298 ;;   (f) ; call stuff
299 ;;                                      ; restore preserved registers
300 ;;   ;; (for-each (lambda (i)
301 ;;   ;;               (emit (string-append
302 ;;   ;;                      "pop %r"
303 ;;   ;;                      (number->string i))))
304 ;;   ;;             '(15 14 13 12))
305 ;;   ;; (emit "pop %rbx")
306 ;;   (emit "pop %rbp")
307 ;;   (emit "ret"))
308
309                                         ; 24(%rbp) mem arg 1
310                                         ; 16(%rbp) mem arg 0          prev frame
311                                         ; -----------------------
312                                         ;  8(%rbp) return address     cur frame
313                                         ;  0(%rbp) prev %rbp
314                                         ; -8(%rbp) do what you want
315                                         ;  ...     do what you want
316                                         ;  0(%rsp) do what you want
317
318 (define (param-register n)
319   (case n
320     (0 "%rdi")
321     (1 "%rsi")
322     (2 "%rdx")
323     (3 "%rcx")
324     (4 "%r8")
325     (5 "%r9")
326     (else (error #f "need to test out the below"))
327     (else (format "~a(%rsp)" (- n 6)))))
328
329 (define (codegen program)
330   (let* ((extract-res-0 (extract-strings program))
331          (strings (car extract-res-0))
332          (extract-res-1 (extract-lambdas (cdr extract-res-0)))
333          (lambdas (car extract-res-1))
334          (xform-prog (cdr extract-res-1)))
335
336     (emit "\t.global _start")
337     (emit "\t.text")
338 ;    (emit ".p2align 4,,15") is this needed?
339
340     (for-each codegen-lambda lambdas)
341
342     (emit "_start:")
343     (emit "movq %rsp, %rbp")            ; set up the base pointer
344     (codegen-expr xform-prog 0 '())
345
346                                         ; exit syscall
347     (emit "mov %rax, %rdi")
348     (emit "mov $60, %rax")
349     (emit "syscall")
350
351     (emit "\t.data")
352
353     (for-each codegen-string-data strings)))
354
355 (define (compile-to-binary program output)
356   (when (not (eq? (typecheck program) 'int)) (error #f "not an int"))
357   (let ([tmp-path "/tmp/a.s"])
358     (when (file-exists? tmp-path) (delete-file tmp-path))
359     (with-output-to-file tmp-path
360       (lambda () (codegen program)))
361     (system (format "clang -nostdlib /tmp/a.s -o ~a" output))))