Add type annotation pass.

1 files changed, 376 insertions, 5 deletions

diff --git a/chickadee/graphics/seagull.scm b/chickadee/graphics/seagull.scm
index df5b9d1..7eb0b14 100644
--- a/chickadee/graphics/seagull.scm
+++ b/chickadee/graphics/seagull.scm
@@ -65,6 +65,14 @@
 (define (primitive-call? x)
   (memq x '(+ - * / = < <= > >=)))
 
+(define (difference a b)
+  (match a
+    (() b)
+    ((x . rest)
+     (if (memq x b)
+         (difference rest (delq x b))
+         (cons x (difference rest b))))))
+
 
 ;;;
 ;;; Lexical environments
@@ -76,6 +84,9 @@
 (define (empty-env)
   '())
 
+(define-syntax-rule (env (key value) ...)
+  (list (cons key value) ...))
+
 (define &seagull-unbound-variable-error
   (make-exception-type '&seagull-unbound-variable-error &error '(name)))
 
@@ -113,6 +124,14 @@
           (proc name exp)))
        env))
 
+(define (env-fold proc init env)
+  (fold (lambda (e memo)
+          (match e
+            ((name . exp)
+             (proc name exp memo))))
+        init
+        env))
+
 (define (env-for-each proc env)
   (for-each (match-lambda
               ((name . exp)
@@ -363,22 +382,374 @@
 
 
 ;;;
-;;; Explicit tail call marking
+;;; Recursion analysis
 ;;;
 
-;; Make note of which function calls are in tail position.  GLSL does
-;; not allow for recursive function calls, but Seagull allows them in
-;; tail position only because they can be compiled to a loop in GLSL.
+;; GLSL does not allow for recursive function calls, but Seagull
+;; allows the current function to call itself in tail position because
+;; they can be compiled to a loop in GLSL.
 
 
 ;;;
-;;; Type inference
+;;; Type inference: Typed expressions
 ;;;
 
 ;; Transform all program expressions into typed expressions by
 ;; applying a variant of the Hindley-Milner type inference algorithm.
 ;; GLSL is a statically typed language, but thanks to type inference
 ;; Seagull code doesn't need type annotations.
+;;
+;; This section includes all the kinds of types an expression may have
+;; and various transformation procedures to manipulate them.
+
+;; Primitive types:
+(define (primitive-type name)
+  `(primitive ,name))
+
+(define (primitive-type? obj)
+  (match obj
+    (('primitive _) #t)
+    (_ #f)))
+
+(define int-type (primitive-type 'int))
+(define float-type (primitive-type 'float))
+(define bool-type (primitive-type 'bool))
+
+;; Type variables:
+(define unique-type-variable-counter (make-parameter 0))
+
+(define (unique-type-variable-number)
+  (let ((n (unique-type-variable-counter)))
+    (unique-type-variable-counter (+ n 1))
+    n))
+
+(define (unique-type-variable-name)
+  (string->symbol
+   (format #f "T~a" (unique-type-variable-number))))
+
+(define (fresh-type-variable)
+  `(tvar ,(unique-type-variable-name)))
+
+(define (type-variable? obj)
+  (match obj
+    (('tvar _) #t)
+    (_ #f)))
+
+;; Function types:
+(define (function-type parameters returns)
+  `(-> ,parameters ,returns))
+
+(define (function-type? obj)
+  (match obj
+    (('-> _ _) #t)
+    (_ #f)))
+
+(define (function-type-parameters type)
+  (match type
+    (('-> params _) params)))
+
+(define (function-type-returns type)
+  (match type
+    (('-> _ returns) returns)))
+
+;; For all types:
+(define (for-all-type quantifiers type)
+  `(for-all ,quantifiers ,type))
+
+(define (for-all-type? obj)
+  (match obj
+    (('for-all _ _) #t)
+    (_ #f)))
+
+(define (for-all-type-quantifiers type)
+  (match type
+    (('for-all q _) q)))
+
+(define (for-all-type-ref type)
+  (match type
+    (('for-all _ t) t)))
+
+;; Intersection types:
+(define (intersection-type . types)
+  `(^ ,types))
+
+(define (intersection-type? obj)
+  (match obj
+    (('^ _) #t)
+    (_ #f)))
+
+(define (intersection-type-ref type)
+  (match type
+    (('^ types) types)))
+
+(define (type? obj)
+  (or (primitive-type? obj)
+      (type-variable? obj)
+      (function-type? obj)
+      (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))))))
+
+(define (occurs? a b)
+  (cond
+   ((and (type-variable? a) (type-variable? b))
+    (eq? a b))
+   ((and (type-variable? a) (function-type? b))
+    (or (occurs? a (function-type-parameters b))
+        (occurs? a (function-type-returns b))))
+   ((and (type? a) (list? b))
+    (any (lambda (b*) (occurs? a b*)) b))
+   (else #f)))
+
+(define (apply-substitution-to-type type from to)
+  (cond
+   ((primitive-type? type) type)
+   ((type-variable? type)
+    (if (eq? type from) to type))
+   ((function-type? type)
+    (function-type
+     (map (lambda (param-type)
+            (apply-substitution-to-type param-type from to))
+          (function-type-parameters type))
+     (map (lambda (return-type)
+            (apply-substitution-to-type return-type from to))
+          (function-type-returns type))))
+   (else (error "invalid type" type))))
+
+(define (apply-substitutions-to-type type env)
+  (env-fold (lambda (from to type*)
+              (apply-substitution-to-type type* from to))
+            type
+            env))
+
+(define (apply-substitution-to-env from to env)
+  (env-fold (lambda (a b env*)
+              (extend-env (if (eq? a from) to a)
+                          (if (eq? b from) to b)
+                          env*))
+            (empty-env)
+            env))
+
+(define (substitute-type from to env)
+  (let* ((to* (apply-substitutions-to-type to env)))
+    (and (not (occurs? from to*))
+         (let ((env* (apply-substitution-to-env from to* env)))
+           (extend-env from to env*)))))
+
+(define (free-variables-in-type type)
+  (cond
+   ((primitive-type? type) '())
+   ((type-variable? type) 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)))))
+   ((for-all-type? type)
+    (fold delete
+          (free-variables-in-type (for-all-type-ref type))
+          (for-all-type-quantifiers type)))
+   (else (error "unknown type" type))))
+
+(define (free-variables-in-for-all for-all)
+  (difference (for-all-type-quantifiers for-all)
+              (free-variables-in-type (for-all-type-ref for-all))))
+
+(define (free-variables-in-env env)
+  (delete-duplicates
+   (env-fold (lambda (_name type vars)
+               (cond
+                ((type-variable? type)
+                 (cons (free-variables-in-type type)
+                       vars))
+                ((for-all-type? type)
+                 (cons (free-variables-in-for-all type)
+                       vars))
+                (else vars)))
+             '()
+             env)))
+
+(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.
+
+(define (fresh-type-variables-for-list lst)
+  (map (lambda (_x) (fresh-type-variable)) lst))
+
+(define (annotate:list exps env)
+  (map (lambda (exp) (annotate-exp exp env)) exps))
+
+(define (annotate:if predicate consequent alternate env)
+  (texp (list (fresh-type-variable))
+        `(if ,(annotate-exp predicate env)
+             ,(annotate-exp consequent env)
+             ,(annotate-exp alternate env))))
+
+(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 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*)))
+
+(define (annotate:lambda params body env)
+  ;; Each function parameter gets a fresh type variable.
+  (define param-types (fresh-type-variables-for-list params))
+  ;; The type environment is extended with the function parameters.
+  (define env* (fold extend-env env params param-types))
+  ;; The body is annotated in the new environment.
+  (define body* (annotate-exp body env*))
+  (texp (list (function-type param-types (texp-types body*)))
+        `(lambda ,params ,body*)))
+
+(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 args* (annotate-list args env))
+  (texp (fresh-type-variables-for-list (list operator*))
+        `(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*)))
+
+(define* (annotate:top-level bindings body env #:optional (result '()))
+  (match bindings
+    (()
+     `(top-level ,result ,(annotate-exp body env)))
+    (((name exp) . rest)
+     (define exp*
+       (let ((x (annotate-exp exp env)))
+         (texp (list (generalize (single-type x) env))
+               (texp-exp x))))
+     (define env* (extend-env name (texp-types exp*) env))
+     (define result* (cons (list name exp*) result))
+     (annotate-top-level rest body env* result*))))
+
+(define (annotate-exp exp env)
+  (match exp
+    ((? exact-integer?)
+     (texp (list int-type) exp))
+    ((? float?)
+     (texp (list float-type) exp))
+    ((? boolean?)
+     (texp (list bool-type) exp))
+    (('var var _)
+     (texp (lookup var env) exp))
+    (('if predicate consequent alternate)
+     (annotate:if predicate consequent alternate env))
+    (('let ((names exps) ...) body)
+     (annotate:let names exps body env))
+    (('lambda (params ...) body)
+     (annotate:lambda params body env))
+    (('primcall operator args ...)
+     (annotate:primitive-call operator args env))
+    (('call operator args ...)
+     (annotate:call operator args env))
+    (('top-level bindings body)
+     (annotate:top-level bindings body env))))
+
+(define (annotate-exp* exp)
+  (parameterize ((unique-type-variable-counter 0))
+    (annotate-exp exp (top-level-type-env))))
+
+
+;;;
+;;; Type inference: Constraints
+;;;
+
+;; 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)
+  (match exp
+    (('t ((? type? types) ...) ('if predicate consequent alternate))
+     (constrain:if types predicate consequent alternate))))
+
+
+;;;
+;;; Type inference: Unification
+;;;
+
+;; Solve type constraints.
+
+
+;;;
+;;; Type inference: Resolution
+;;;
+
+;; Replace the variables in the original typed expression to form a
+;; fully typed expression that is ready for emission to GLSL.
 
 
 ;;;