WIP on typechecker refactor
[scheme.git] / typecheck.scm
1 (load "ast.scm")
2
3 (define (abs? t)
4   (and (list? t) (eq? (car t) 'abs)))
5
6 (define (tvar? t)
7   (and (not (list? t)) (not (concrete? t)) (symbol? t)))
8
9 (define (concrete? t)
10   (case t
11     ('int #t)
12     ('bool #t)
13     ('void #t)
14     (else #f)))
15
16 (define (pretty-type t)
17   (cond ((abs? t)
18          (string-append
19           (if (abs? (cadr t))
20               (string-append "(" (pretty-type (cadr t)) ")")
21               (pretty-type (cadr t)))
22           " -> "
23           (pretty-type (caddr t))))
24         (else (symbol->string t))))
25
26 (define (pretty-constraints cs)
27   (string-append "{"
28                  (fold-left string-append
29                             ""
30                             (map (lambda (c)
31                                    (string-append
32                                     (pretty-type (car c))
33                                     ": "
34                                     (pretty-type (cdr c))
35                                     ", "))
36                                  cs))
37                  "}"))
38
39                                         ; ('a, ('b, 'a))
40 (define (env-lookup env n)
41   (if (null? env) (error #f "empty env")                        ; it's a type equality
42       (if (eq? (caar env) n)
43           (cdar env)
44           (env-lookup (cdr env) n))))
45
46 (define (env-insert env n t)
47   (cons (cons n t) env))
48
49 (define abs-arg cadr)
50
51 (define cur-tvar 0)
52 (define (fresh-tvar)
53   (begin
54     (set! cur-tvar (+ cur-tvar 1))
55     (string->symbol
56      (string-append "t" (number->string (- cur-tvar 1))))))
57
58 (define (last xs)
59   (if (null? (cdr xs))
60       (car xs)
61       (last (cdr xs))))
62
63 (define (normalize prog) ; (+ a b) -> ((+ a) b)
64   (case (ast-type prog)
65     ('lambda 
66                                         ; (lambda (x y) (+ x y)) -> (lambda (x) (lambda (y) (+ x y)))
67         (if (> (length (lambda-args prog)) 1)
68             (list 'lambda (list (car (lambda-args prog)))
69                   (normalize (list 'lambda (cdr (lambda-args prog)) (caddr prog))))
70             (list 'lambda (lambda-args prog) (normalize (caddr prog)))))
71     ('app
72      (if (null? (cddr prog))
73          `(,(normalize (car prog)) ,(normalize (cadr prog))) ; (f a)
74          (normalize `(,(list (normalize (car prog)) (normalize (cadr prog)))
75                       ,@(cddr prog))))) ; (f a b)
76     ('let
77         (append (list 'let
78                       (map (lambda (x) `(,(car x) ,(normalize (cadr x))))
79                            (let-bindings prog)))
80                 (map normalize (let-body prog))))
81     (else (ast-traverse normalize prog))))
82
83 (define (builtin-type x)
84   (case x
85     ('+ '(abs int (abs int int)))
86     ('- '(abs int (abs int int)))
87     ('* '(abs int (abs int int)))
88     ('! '(abs bool bool))
89     ('= '(abs int (abs int bool)))
90     ('bool->int '(abs bool int))
91     ('print '(abs string void))
92     (else #f)))
93
94 (define (check env x)
95   (display "check: ")
96   (display x)
97   (display "\n\t")
98   (display env)
99   (newline)
100   (let
101       ((res
102         (case (ast-type x)
103           ('int-literal (list '() 'int))
104           ('bool-literal (list '() 'bool))
105           ('string-literal (list '() 'string))
106           ('builtin (list '() (builtin-type x)))
107
108           ('if
109            (let* ((cond-type-res (check env (cadr x)))
110                   (then-type-res (check env (caddr x)))
111                   (else-type-res (check env (cadddr x)))
112                   (then-eq-else-cs (~ (cadr then-type-res)
113                                       (cadr else-type-res)))
114                   (cs (constraint-merge
115                        (car then-type-res)
116                        (constraint-merge (car else-type-res)
117                                          then-eq-else-cs)))
118                   (return-type (substitute cs (cadr then-type-res))))
119              (when (not (eqv? (cadr cond-type-res) 'bool))
120                (error #f "if condition isn't bool"))
121              (list cs return-type)))
122           
123           ('var (list '() (env-lookup env x)))
124           ('let
125                                         ; takes in the current environment and a scc
126                                         ; returns new environment with scc's types added in
127               (let* ([components (reverse (sccs (graph (let-bindings x))))]
128                      [process-component
129                       (lambda (acc comps)
130                         (let*
131                                         ; create a new env with tvars for each component
132                                         ; e.g. scc of (x y)
133                                         ; scc-env = ((x . t0) (y . t1))
134                             ([scc-env
135                               (fold-left
136                                (lambda (acc c)
137                                  (env-insert acc c (fresh-tvar)))
138                                acc comps)]
139                                         ; typecheck each component
140                              [type-results
141                               (map
142                                (lambda (c)
143                                  (let ([body (cadr (assoc c (let-bindings x)))])
144                                    (check scc-env body)))
145                                comps)]
146                                         ; collect all the constraints in the scc
147                              [cs
148                               (fold-left
149                                (lambda (acc res c)
150                                  (constraint-merge
151                                   (constraint-merge
152                                         ; unify with tvars from scc-env
153                                         ; result ~ tvar
154                                    (~ (env-lookup scc-env c) (cadr res))
155                                    (car res))                             
156                                   acc))
157                                '() type-results comps)]
158                                         ; substitute *only* the bindings in this scc
159                              [new-env
160                               (map (lambda (x)
161                                      (if (memv (car x) comps)
162                                          (cons (car x) (substitute cs (cdr x)))
163                                          x))
164                                    scc-env)])
165                                         (display "cs:")
166                 (display cs)
167                 (newline)
168                           new-env))]
169                      [new-env (fold-left process-component env components)])
170                 (check new-env (last (let-body x)))))
171           
172           ('lambda
173               (let* [(new-env (env-insert env (lambda-arg x) (fresh-tvar)))
174
175                      (body-type-res (check new-env (lambda-body x)))
176                      (cs (car body-type-res))
177                      (subd-env (substitute-env (car body-type-res) new-env))
178                      (arg-type (env-lookup subd-env (lambda-arg x)))
179                      (resolved-arg-type (substitute cs arg-type))]
180                 ;; (display "lambda:\n\t")
181                 ;; (display prog)
182                 ;; (display "\n\t")
183                 ;; (display cs)
184                 ;; (display "\n\t")
185                 ;; (display (format "subd-env: ~a\n" subd-env))
186                 ;; (display resolved-arg-type)
187                 ;; (newline)
188                 (list (car body-type-res)
189                       (list 'abs
190                             resolved-arg-type
191                             (cadr body-type-res)))))
192           
193           ('app ; (f a)
194            (if (eqv? (car x) (cadr x))
195                                         ; recursive function (f f)
196                (let* [(func-type (env-lookup env (car x)))
197                       (return-type (fresh-tvar))
198                       (other-func-type `(abs ,func-type ,return-type))
199                       (cs (~ func-type other-func-type))
200                       (resolved-return-type (substitute cs return-type))]
201                  (list cs resolved-return-type))
202
203                                         ; regular function
204                (let* ((arg-type-res (check env (cadr x)))
205                       (arg-type (cadr arg-type-res))
206                       (func-type-res (check env (car x)))
207                       (func-type (cadr func-type-res))
208                       
209                                         ; f ~ a -> t0
210                       (func-c (~
211                                (substitute (car arg-type-res) func-type)
212                                `(abs ,arg-type ,(fresh-tvar))))
213                       (cs (constraint-merge
214                            (constraint-merge func-c (car arg-type-res))
215                            (car func-type-res)))
216                       
217                       (resolved-func-type (substitute cs func-type))
218                       (resolved-return-type (caddr resolved-func-type)))
219                  ;; (display "app:\n")
220                  ;; (display cs)
221                  ;; (display "\n")
222                  ;; (display func-type)
223                  ;; (display "\n")
224                  ;; (display resolved-func-type)
225                  ;; (display "\n")
226                  ;; (display arg-type-res)
227                  ;; (display "\n")
228                  (if (abs? resolved-func-type)
229                      (let ((return-type (substitute cs (caddr resolved-func-type))))
230                        (list cs return-type))
231                      (error #f "not a function"))))))))
232     (display "result of ")
233     (display x)
234     (display ":\n\t")
235     (display (pretty-type (cadr res)))
236     (display "\n\t[")
237     (display (pretty-constraints (car res)))
238     (display "]\n")
239     res))
240
241                                         ; we typecheck the lambda calculus only (only single arg lambdas)
242 (define (typecheck prog)
243   (cadr (check '() (normalize prog))))
244
245                                         ; returns a list of constraints
246 (define (~ a b)
247   (let ([res (unify? a b)])
248     (if res
249         res
250         (error #f
251                (format "couldn't unify ~a ~~ ~a" a b)))))
252
253 (define (unify? a b)
254   (cond [(eq? a b) '()]
255         [(tvar? a) (list (cons a b))]
256         [(tvar? b) (list (cons b a))]
257         [(and (abs? a) (abs? b))
258          (let* [(arg-cs (unify? (cadr a) (cadr b)))
259                 (body-cs (unify? (substitute arg-cs (caddr a))
260                                  (substitute arg-cs (caddr b))))]
261            (constraint-merge body-cs arg-cs))]
262         [else #f]))
263
264 (define (substitute cs t)
265   (cond
266    [(tvar? t)
267     (if (assoc t cs)
268         (cdr (assoc t cs))
269         t)]
270    [(abs? t) `(abs ,(substitute cs (cadr t))
271                    ,(substitute cs (caddr t)))]
272    [else t]))
273
274                                         ; applies substitutions to all variables in environment
275 (define (substitute-env cs env)
276   (map (lambda (x) (cons (car x) (substitute cs (cdr x)))) env))
277
278                                         ; composes constraints a onto b and merges, i.e. applies a to b
279                                         ; a should be the "more important" constraints
280 (define (constraint-merge a b)
281   (define (f constraint)
282     (cons (car constraint)
283           (substitute a (cdr constraint))))
284   
285   (define (most-concrete a b)
286     (cond
287      [(tvar? a) b]
288      [(tvar? b) a]
289      [(and (abs? a) (abs? b))
290       `(abs ,(most-concrete (cadr a) (cadr b))
291             ,(most-concrete (caddr a) (caddr b)))]
292      [(abs? a) b]
293      [(abs? b) a]
294      [else (error #f "impossible! most-concrete")]))
295
296   (define (union p q)
297     (cond
298      [(null? p) q]
299      [(null? q) p]
300      [else
301       (let ([x (car q)])
302         (if (assoc (car x) p)
303             (if (eqv? (most-concrete (cddr (assoc (car x) p))
304                                      (cdr x))
305                       (cdr x))
306                 (cons x (union (filter (p) (not (eqv? 
307         
308   
309   (define (union p q)
310     (append (filter (lambda (x) (not (assoc (car x) p)))
311                     q)
312             p))
313   (union a (map f b)))
314
315
316 ;;                                      ; a1 -> a2 ~ a3 -> a4;
317 ;;                                      ; a1 -> a2 !~ bool -> bool
318 ;;                                      ; basically can the tvars be renamed
319 (define (types-equal? x y)
320   (let ([cs (unify? x y)])
321     (if (not cs) #f     
322         (let*
323             ([test (lambda (acc c)
324                      (and acc
325                           (tvar? (car c)) ; the only substitutions allowed are tvar -> tvar
326                           (tvar? (cdr c))))])
327           (fold-left test #t cs)))))
328
329                                         ; input: a list of binds ((x . y) (y . 3))
330                                         ; returns: pair of verts, edges ((x y) . (x . y))
331 (define (graph bs)
332   (define (go bs orig-bs)
333     (define (find-refs prog)
334       (ast-collect
335        (lambda (x)
336          (case (ast-type x)
337                                         ; only count a reference if its a binding
338            ['var (if (assoc x orig-bs) (list x) '())]
339            [else '()]))
340        prog))
341     (if (null? bs)
342         '(() . ())
343         (let* [(bind (car bs))
344
345                (vert (car bind))
346                (refs (find-refs (cdr bind)))
347                (edges (map (lambda (x) (cons vert x))
348                            refs))
349
350                (rest (if (null? (cdr bs))
351                          (cons '() '())
352                          (go (cdr bs) orig-bs)))
353                (total-verts (cons vert (car rest)))
354                (total-edges (append edges (cdr rest)))]
355           (cons total-verts total-edges))))
356   (go bs bs))
357
358 (define (successors graph v)
359   (define (go v E)
360     (if (null? E)
361         '()
362         (if (eqv? v (caar E))
363             (cons (cdar E) (go v (cdr E)))
364             (go v (cdr E)))))
365   (go v (cdr graph)))
366
367                                         ; takes in a graph (pair of vertices, edges)
368                                         ; returns a list of strongly connected components
369
370                                         ; ((x y w) . ((x . y) (x . w) (w . x))
371
372                                         ; =>
373                                         ; .->x->y
374                                         ; |  |
375                                         ; |  v
376                                         ; .--w
377
378                                         ; ((x w) (y))
379
380                                         ; this uses tarjan's algorithm, to get reverse
381                                         ; topological sorting for free
382 (define (sccs graph)
383   
384   (let* ([indices (make-hash-table)]
385          [lowlinks (make-hash-table)]
386          [on-stack (make-hash-table)]
387          [current 0]
388          [stack '()]
389          [result '()])
390
391     (define (index v)
392       (get-hash-table indices v #f))
393     (define (lowlink v)
394       (get-hash-table lowlinks v #f))
395
396     (letrec
397         ([strong-connect
398           (lambda (v)
399             (begin
400               (put-hash-table! indices v current)
401               (put-hash-table! lowlinks v current)
402               (set! current (+ current 1))
403               (push! stack v)
404               (put-hash-table! on-stack v #t)
405
406               (for-each
407                (lambda (w)
408                  (if (not (hashtable-contains? indices w))
409                                         ; successor w has not been visited, recurse
410                      (begin
411                        (strong-connect w)
412                        (put-hash-table! lowlinks
413                                         v
414                                         (min (lowlink v) (lowlink w))))
415                                         ; successor w has been visited
416                      (when (get-hash-table on-stack w #f)
417                        (put-hash-table! lowlinks v (min (lowlink v) (index w))))))
418                (successors graph v))
419
420               (when (= (index v) (lowlink v))
421                 (let ([scc
422                        (let new-scc ()
423                          (let ([w (pop! stack)])
424                            (put-hash-table! on-stack w #f)
425                            (if (eqv? w v)
426                                (list w)
427                                (cons w (new-scc)))))])
428                   (set! result (cons scc result))))))])
429       (for-each
430        (lambda (v)
431          (when (not (hashtable-contains? indices v)) ; v.index == -1
432            (strong-connect v)))
433        (car graph)))
434     result))
435