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