Add unification pass.

1 files changed, 268 insertions, 89 deletions

diff --git a/chickadee/graphics/seagull.scm b/chickadee/graphics/seagull.scm
index 7eb0b14..c53e034 100644
--- a/chickadee/graphics/seagull.scm
+++ b/chickadee/graphics/seagull.scm
@@ -84,7 +84,7 @@
 (define (empty-env)
   '())
 
-(define-syntax-rule (env (key value) ...)
+(define-syntax-rule (new-env (key value) ...)
   (list (cons key value) ...))
 
 (define &seagull-unbound-variable-error
@@ -118,6 +118,9 @@
 (define (env-names env)
   (map car env))
 
+(define (env-values env)
+  (map cdr env))
+
 (define (env-map proc env)
   (map (match-lambda
          ((name . exp)
@@ -188,8 +191,9 @@
 (define (expand-let names exps body env)
   (if (null? names)
       (expand body env)
-      (let* ((env* (compose-envs (alpha-convert names) env))
-             (bindings* (map list (lookup-all names env*) exps)))
+      (let* ((exps* (map (lambda (exp) (expand exp env)) exps))
+             (env* (compose-envs (alpha-convert names) env))
+             (bindings* (map list (lookup-all names env*) exps*)))
         `(let ,bindings* ,(expand body env*)))))
 
 (define (expand-let* bindings body env)
@@ -414,6 +418,11 @@
 (define int-type (primitive-type 'int))
 (define float-type (primitive-type 'float))
 (define bool-type (primitive-type 'bool))
+(define vec2-type (primitive-type 'vec2))
+(define vec3-type (primitive-type 'vec3))
+(define vec4-type (primitive-type 'vec4))
+(define mat3-type (primitive-type 'mat3))
+(define mat4-type (primitive-type 'mat4))
 
 ;; Type variables:
 (define unique-type-variable-counter (make-parameter 0))
@@ -489,11 +498,11 @@
       (intersection-type? obj)))
 
 (define (top-level-type-env)
-  `((+ . ,(intersection-type
-           (function-type (list int-type int-type)
-                          (list int-type))
-           (function-type (list float-type float-type)
-                          (list float-type))))))
+  `((+ . ,(list (intersection-type
+                 (function-type (list int-type int-type)
+                                (list int-type))
+                 (function-type (list float-type float-type)
+                                (list float-type)))))))
 
 (define (occurs? a b)
   (cond
@@ -510,7 +519,7 @@
   (cond
    ((primitive-type? type) type)
    ((type-variable? type)
-    (if (eq? type from) to type))
+    (if (equal? type from) to type))
    ((function-type? type)
     (function-type
      (map (lambda (param-type)
@@ -541,16 +550,53 @@
          (let ((env* (apply-substitution-to-env from to* env)))
            (extend-env from to env*)))))
 
-(define (free-variables-in-type type)
+;; Typed expressions:
+(define (texp types exp)
+  `(t ,types ,exp))
+
+(define (texp? obj)
+  (match obj
+    (('t _ _) #t)
+    (_ #f)))
+
+(define (texp-types texp)
+  (match texp
+    (('t types _) types)))
+
+(define (texp-exp texp)
+  (match texp
+    (('t _ exp) exp)))
+
+(define (single-type texp)
+  (match (texp-types texp)
+    ((type) type)
+    (_ (error "expected only 1 type" texp))))
+
+
+;;;
+;;; Type inference: Annotation
+;;;
+
+;; Convert untyped Seagull expressions into typed expressions with
+;; variable components that need to be solved for.
+
+(define (lookup-type name env)
+  (pk 'lookup-type name env)
+  (let ((type (pk 'looked-up (lookup name env))))
+    (if (for-all-type? type)
+        (instantiate type)
+        type)))
+
+(define* (free-variables-in-type type)
   (cond
    ((primitive-type? type) '())
-   ((type-variable? type) type)
+   ((type-variable? type) (list type))
    ((function-type? type)
-    (delete-duplicates
-     (append (map free-variables-in-type
-                  (function-type-parameters type))
-             (map free-variables-in-type
-                  (function-type-returns type)))))
+    (let ((params (function-type-parameters type)))
+      (filter (lambda (t) (member t params))
+              (delete-duplicates
+               (append-map free-variables-in-type
+                           (function-type-returns type))))))
    ((for-all-type? type)
     (fold delete
           (free-variables-in-type (for-all-type-ref type))
@@ -575,50 +621,24 @@
              '()
              env)))
 
+(define (generalize type env)
+  (if (function-type? type)
+      ;; (let ((quantifiers (difference (free-variables-in-type type)
+      ;;                                (free-variables-in-env env))))
+      ;;   (for-all-type quantifiers type))
+      (match (difference (free-variables-in-type type)
+                         (free-variables-in-env env))
+        (() type)
+        ((quantifiers) (for-all-type quantifiers type)))
+      type))
+
 (define (instantiate for-all)
   (define subs
     (fold (lambda (var env)
             (extend-env var (fresh-type-variable) env))
           (empty-env)
           (for-all-type-quantifiers for-all)))
-  (apply-substitutions-to-type (for-all-type-ref for-all) subs))
-
-(define (generalize type env)
-  (match (difference (free-variables-in-type type)
-                     (free-variables-in-env env))
-    (() type)
-    (quantifiers
-     (for-all-type quantifiers type))))
-
-;; Typed expressions:
-(define (texp types exp)
-  `(t ,types ,exp))
-
-(define (texp? obj)
-  (match obj
-    (('t _ _) #t)
-    (_ #f)))
-
-(define (texp-types texp)
-  (match texp
-    (('t types _) types)))
-
-(define (texp-exp texp)
-  (match texp
-    (('t _ exp) exp)))
-
-(define (single-type texp)
-  (match (texp-types texp)
-    ((type) type)
-    (_ (error "expected only 1 type" texp))))
-
-
-;;;
-;;; Type inference: Annotation
-;;;
-
-;; Convert untyped Seagull expressions into typed expressions with
-;; variable components that need to be solved for.
+  (pk 'instantiated (apply-substitutions-to-type (for-all-type-ref for-all) subs)))
 
 (define (fresh-type-variables-for-list lst)
   (map (lambda (_x) (fresh-type-variable)) lst))
@@ -634,23 +654,17 @@
 
 (define (annotate:let names exps body env)
   (define exps* (annotate-list exps env))
-  ;; Every expression type that is being bound in a 'let' must only
-  ;; return a single value.  Function types must be generalized so
-  ;; that functions like (lambda (x) x) can truly be applied to any
-  ;; type.
-  (define exp-types (map (lambda (exp)
-                           (generalize (single-type exp) env))
-                         exps*))
-  ;; The annotated expressions are rewritten to use generalized types.
-  (define exps** (map (lambda (exp type)
-                        (texp (list type) (texp-exp exp)))
-                      exps* exp-types))
-  ;; The type environment is extended with the generalized bindings.
+  (define exp-types (map single-type exps*))
   (define env* (fold extend-env env names exp-types))
-  ;; The 'let' body is annotated in the new environment.
   (define body* (annotate-exp body env*))
-  (texp (texp-types body*)
-        `(let ,(map list names exps**) ,body*)))
+  (texp (list (function-type (fresh-type-variables-for-list exps)
+                             (texp-types body*)))
+        `(let ,(map list names exps*) ,body*)))
+
+;; (let ((a 1) (b 2)) (+ a b))
+;; ((lambda (a b) (+ a b)) 1 2)
+;; (-> (T0 T1) (int))
+;; (-> (int int) (int))
 
 (define (annotate:lambda params body env)
   ;; Each function parameter gets a fresh type variable.
@@ -665,29 +679,33 @@
 (define (annotate:primitive-call operator args env)
   ;; The type signature of primitive functions can be looked up
   ;; directly in the environment.
-  (define operator* (lookup operator env))
+  (define operator-types (lookup-type operator env))
+  (define operator* (texp operator-types operator))
   (define args* (annotate-list args env))
   (texp (fresh-type-variables-for-list (list operator*))
-        `(primcall ,operator ,@args*)))
+        `(primcall ,operator* ,@args*)))
 
 (define (annotate:call operator args env)
   (define operator* (annotate-exp operator env))
   (define args* (annotate-list args env))
   (texp (fresh-type-variables-for-list (texp-types operator*))
-        `(call operator* ,@args*)))
+        `(call ,operator* ,@args*)))
 
 (define* (annotate:top-level bindings body env #:optional (result '()))
   (match bindings
     (()
-     `(top-level ,result ,(annotate-exp body env)))
+     (let ((body* (pk 'body (annotate-exp body env))))
+       (texp (texp-types body*) `(top-level ,result ,body*))))
     (((name exp) . rest)
      (define exp*
        (let ((x (annotate-exp exp env)))
+         ;; Function types must be generalized so that functions like
+         ;; (lambda (x) x) can truly be applied to any type.
          (texp (list (generalize (single-type x) env))
                (texp-exp x))))
-     (define env* (extend-env name (texp-types exp*) env))
+     (define env* (extend-env name (single-type exp*) env))
      (define result* (cons (list name exp*) result))
-     (annotate-top-level rest body env* result*))))
+     (annotate:top-level rest body env* result*))))
 
 (define (annotate-exp exp env)
   (match exp
@@ -698,7 +716,7 @@
     ((? boolean?)
      (texp (list bool-type) exp))
     (('var var _)
-     (texp (lookup var env) exp))
+     (texp (list (lookup-type var env)) exp))
     (('if predicate consequent alternate)
      (annotate:if predicate consequent alternate env))
     (('let ((names exps) ...) body)
@@ -724,17 +742,61 @@
 ;; Generate type constraints for a program.
 
 (define (constrain:if types predicate consequent alternate)
-  (append (env ((list bool-type) (texp-types predicate))
-               (types (texp-types consequent))
-               (types (texp-types alternate)))
-          (constrain-exp predicate)
-          (constrain-exp consequent)
-          (constrain-exp alternate)))
-
-(define (constrain-exp exp)
+  (append (new-env ((list bool-type) (texp-types predicate))
+                   (types (texp-types consequent))
+                   (types (texp-types alternate)))
+          (constrain-texp predicate)
+          (constrain-texp consequent)
+          (constrain-texp alternate)))
+
+(define (constrain:let type rhs body)
+  (append (new-env (type (function-type (map single-type rhs)
+                                        (texp-types body))))
+          (constrain-texp body)
+          (append-map constrain-texp rhs)))
+
+(define (constrain:lambda type body)
+  (append (new-env ((function-type-returns
+                     (if (for-all-type? type)
+                         (for-all-type-ref type)
+                         type))
+                    (texp-types body)))
+          (constrain-texp body)))
+
+(define (constrain:primitive-call types operator args)
+  (append (new-env ((texp-types operator)
+                    (list (function-type (map single-type args) types))))
+          (constrain-texp operator)
+          (append-map constrain-texp args)))
+
+(define (constrain:call types operator args)
+  (append (new-env ((texp-types operator)
+                    (list (function-type (map single-type args) types))))
+        (constrain-texp operator)
+        (append-map constrain-texp args)))
+
+(define (constrain:top-level types rhs body)
+  (append (constrain-texp body)
+          (append-map constrain-texp rhs)))
+
+(define (constrain-texp exp)
   (match exp
-    (('t ((? type? types) ...) ('if predicate consequent alternate))
-     (constrain:if types predicate consequent alternate))))
+    (('t (types ...) ('if predicate consequent alternate))
+     (constrain:if types predicate consequent alternate))
+    (('t (type) ('let ((_ rhs) ...) body))
+     (constrain:let type rhs body))
+    (('t (type) ('lambda _ body))
+     (pk 'constrain:lambda type body)
+     (constrain:lambda type body))
+    (('t (types ...) ('primcall operator args ...))
+     (constrain:primitive-call types operator args))
+    (('t (types ...) ('call operator args ...))
+     (pk 'constrain:call types operator args)
+     (constrain:call types operator args))
+    (('t (types ...) ('top-level ((_ rhs) ...) body))
+     (constrain:top-level types rhs body))
+    (('t (_ ...) _)
+     (empty-env))))
 
 
 ;;;
@@ -743,6 +805,85 @@
 
 ;; Solve type constraints.
 
+(define %unify-prompt-tag (make-prompt-tag 'unify))
+
+(define (call-with-unify-backtrack thunk failure-handler)
+  (call-with-prompt %unify-prompt-tag
+    thunk
+    (lambda (_k) (failure-handler))))
+
+(define (unify-fail)
+  (abort-to-prompt %unify-prompt-tag))
+
+(define (unify:maybe-substitute var other env)
+  (pk 'maybe-sub var other)
+  (cond
+   ;; Tautology: matching 2 vars that are the same.
+   ((and (type-variable? other) (eq? var other))
+    env)
+   ;; Variable has been bound to some other value, recursively follow
+   ;; it and unify.
+   ((assq-ref env var) =>
+    (lambda (type)
+      (unify type other env)))
+   ;; Substitute variable for value.
+   (else
+    (or (pk 'sub var other (substitute var other env))
+        (unify-fail)))))
+
+(define (unify:constants a b env)
+  (pk 'a a)
+  (pk 'b b)
+  (if (eqv? a b) env (unify-fail)))
+
+(define (unify:lists a rest-a b rest-b env)
+  (unify rest-a rest-b (unify a b env)))
+
+(define (unify:variables a b env)
+  (unify:maybe-substitute a b env))
+
+(define (unify:functions a b env)
+  (unify (function-type-returns a)
+         (function-type-returns b)
+         (unify (function-type-parameters a)
+                (function-type-parameters b)
+                env)))
+
+(define (unify:intersection a b env)
+  (match (filter-map (lambda (t)
+                       (call-with-unify-backtrack
+                        (lambda ()
+                          (unify t b env))
+                        (lambda () #f)))
+                     (intersection-type-ref a))
+    ((env) env)
+    (_ (unify-fail))))
+
+(define (unify . args)
+  (match args
+    ((or ((? type-variable? a) b env)
+         (b (? type-variable? a) env))
+     (pk 'unify-vars a b)
+     (unify:variables a b env))
+    (((? function-type? a) (? function-type? b) env)
+     (pk 'unify-funcs a b)
+     (unify:functions a b env))
+    (((? intersection-type? a) b env)
+     (pk 'unify-intersection a b)
+     (unify:intersection a b env))
+    (((a . rest-a) (b . rest-b) env)
+     (pk 'unify-lists a rest-a b rest-b)
+     (unify:lists a rest-a b rest-b env))
+    ((a b env)
+     (pk 'unify-constants a b)
+     (unify:constants a b env))))
+
+(define (unify-constraints env)
+  (call-with-unify-backtrack
+   (lambda ()
+     (unify (pk 'as (map car env)) (pk 'bs (map cdr env)) '()))
+   (lambda () #f)))
+
 
 ;;;
 ;;; Type inference: Resolution
@@ -751,6 +892,33 @@
 ;; Replace the variables in the original typed expression to form a
 ;; fully typed expression that is ready for emission to GLSL.
 
+(define (resolve:type-variable var env)
+  (let ((type (assq-ref env var)))
+    (if type
+        (resolve type env)
+        var)))
+
+(define (resolve:for-all type env)
+  (let ((quantifiers (filter type-variable?
+                             (resolve (for-all-type-quantifiers type) env)))
+        (ref (resolve (for-all-type-ref type) env)))
+    (if (null? quantifiers) ref (for-all-type quantifiers ref))))
+
+(define (resolve:t))
+
+(define (resolve:list exps env)
+  (map (lambda (exp) (resolve exp env)) exps))
+
+(define (resolve exp env)
+  (match exp
+    ((? type-variable?)
+     (resolve:type-variable exp env))
+    ((? for-all-type?)
+     (resolve:for-all exp env))
+    ((exps ...)
+     (resolve:list exps env))
+    (_ exp)))
+
 
 ;;;
 ;;; GLSL emission
@@ -758,6 +926,7 @@
 
 ;; Transform typed expressions into a string of GLSL code.
 
+
 
 ;;;
 ;;; Compiler front-end
@@ -766,8 +935,18 @@
 ;; Combine all of the compiler passes on a user provided program and
 ;; emit GLSL code if the program is valid.
 
+(define (print-list lst)
+  (for-each (lambda (x) (format #t "~a\n" x)) lst))
+
+;; Substitutions aren't being generated correctly.
 (define (compile-seagull exp)
   (parameterize ((unique-identifier-counter 0))
-    (let* ((expanded (expand exp (top-level-env)))
-           (hoisted (hoist-functions* expanded)))
-      hoisted)))
+    (let* ((expanded (pk 'expanded (expand exp (top-level-env))))
+           (hoisted (pk 'hoisted (hoist-functions* expanded)))
+           (texp (pk 'annotated (annotate-exp* hoisted)))
+           (constraints (pk 'constraints (constrain-texp texp)))
+           (substitutions (pk 'substitutions (unify-constraints constraints))))
+      (list 'annotated texp
+            'constraints constraints
+            'substitutions substitutions
+            'resolved (resolve texp substitutions)))))