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