Add built-in type registry.

1 files changed, 103 insertions, 82 deletions

diff --git a/chickadee/graphics/seagull.scm b/chickadee/graphics/seagull.scm
index dd1cbc6..6ff764e 100644
--- a/chickadee/graphics/seagull.scm
+++ b/chickadee/graphics/seagull.scm
@@ -38,8 +38,9 @@
 ;; - No recursion
 ;;
 ;; TODO:
-;; - Loops
+;; - Seagull unquote
 ;; - User defined structs
+;; - Loops
 ;; - Better error messages (especially around type predicate failure)
 ;; - Helper function modules
 ;; - Shader composition
@@ -219,7 +220,8 @@
 (define-symbolic-type primitive
   primitive-type
   primitive-type?
-  (name primitive-type-name))
+  (name primitive-type-name)
+  (glsl-name primitive-type-glsl-name))
 
 (define-symbolic-type outputs
   outputs-type
@@ -229,14 +231,12 @@
   struct-type
   struct-type?
   (name struct-type-name)
+  (glsl-name struct-type-glsl-name)
   (fields struct-type-fields))
 
 (define (struct-type-ref type field)
   (assq-ref (struct-type-fields type) field))
 
-(define-syntax-rule (define-struct-type (var-name name) (types names) ...)
-  (define var-name (struct-type 'name (list (cons 'names types) ...))))
-
 (define-symbolic-type array
   array-type
   array-type?
@@ -570,10 +570,40 @@
       (values new-pred subs)
       (eval-predicate* new-pred (compose-substitutions subs pred-subs))))
 
+;; Built-in type registry.
+(define *types* (make-hash-table))
+
+(define (lookup-type name)
+  (hashq-ref *types* name))
+
+(define (register-type! name type)
+  (hashq-set! *types* name type))
+
+(define-syntax define-primitive-type
+  (syntax-rules ()
+    ((_ var-name seagull-name)
+     (define-primitive-type var-name
+       seagull-name (symbol->string 'seagull-name)))
+    ((_ var-name seagull-name glsl-name)
+     (begin
+       (define var-name (primitive-type 'seagull-name glsl-name))
+       (register-type! 'seagull-name var-name)))))
+
+(define-syntax define-struct-type
+  (syntax-rules ()
+    ((_ (var-name seagull-name) (types names) ...)
+     (define-struct-type (var-name seagull-name (symbol->string 'seagull-name))
+       (types names) ...))
+    ((_ (var-name seagull-name glsl-name) (types names) ...)
+     (begin
+       (define var-name (struct-type 'seagull-name glsl-name
+                                     (list (cons 'names types) ...)))
+       (register-type! 'seagull-name var-name)))))
+
 ;; Built-in types:
-(define type:int (primitive-type 'int))
-(define type:float (primitive-type 'float))
-(define type:bool (primitive-type 'bool))
+(define-primitive-type type:int int)
+(define-primitive-type type:float float)
+(define-primitive-type type:bool bool)
 (define-struct-type (type:vec2 vec2)
   (type:float x)
   (type:float y))
@@ -588,9 +618,9 @@
   (type:float w))
 ;; TODO: Matrices are technically array types in GLSL, but we are
 ;; choosing to represent them opaquely for now to keep things simple.
-(define type:mat3 (primitive-type 'mat3))
-(define type:mat4 (primitive-type 'mat4))
-(define type:sampler-2d (primitive-type 'sampler2D))
+(define-primitive-type type:mat3 mat3)
+(define-primitive-type type:mat4 mat4)
+(define-primitive-type type:sampler-2d sampler-2d "sampler2D")
 (define type:outputs (outputs-type))
 
 
@@ -656,31 +686,31 @@
   (register-seagull-variable! (make-seagull-variable 'name args ...)))
 
 (define-seagull-variable vertex:position
-  #:glsl-name 'gl_Position
+  #:glsl-name "gl_Position"
   #:type type:vec4
   #:stages '(vertex)
   #:qualifier 'output)
 
 (define-seagull-variable vertex:point-size
-  #:glsl-name 'gl_PointSize
+  #:glsl-name "gl_PointSize"
   #:type type:float
   #:stages '(vertex)
   #:qualifier 'output)
 
 (define-seagull-variable vertex:clip-distance
-  #:glsl-name 'gl_ClipDistance
+  #:glsl-name "gl_ClipDistance"
   #:type type:float
   #:stages '(vertex)
   #:qualifier 'output)
 
 (define-seagull-variable fragment:depth
-  #:glsl-name 'gl_FragDepth
+  #:glsl-name "gl_FragDepth"
   #:type type:float
   #:stages '(fragment)
   #:qualifier 'output)
 
 (define-seagull-variable fragment:coord
-  #:glsl-name 'gl_FragCoord
+  #:glsl-name "gl_FragCoord"
   #:type type:vec4
   #:stages '(fragment)
   #:qualifier 'input)
@@ -1138,7 +1168,7 @@
 (define (unique-identifiers-for-list lst)
   (map (lambda (_x) (unique-identifier)) lst))
 
-(define (top-level-env stage)
+(define (expand:top-level-env stage)
   (fold (lambda (v env)
           (let ((name (seagull-variable-name v)))
             (extend-env name name env)))
@@ -1409,7 +1439,7 @@
      (seagull-syntax-error "unknown form" exp expand))))
 
 (define (expand* exp stage)
-  (expand exp stage (top-level-env stage)))
+  (expand exp stage (expand:top-level-env stage)))
 
 
 ;;;
@@ -2010,6 +2040,18 @@
 (define (type-mismatch a b origin)
   (seagull-type-error "type mismatch" (list a b) origin))
 
+(define (type-descriptor->type desc)
+  (match desc
+    ((? symbol?)
+     (lookup-type desc))
+    (('array desc* (? exact-integer? length) (? exact-integer? rest) ...)
+     (let loop ((rest rest)
+                (prev (array-type (type-descriptor->type desc*) length)))
+       (match rest
+         (() prev)
+         ((length . rest)
+          (loop rest (array-type prev length))))))))
+
 (define (apply-substitution-to-type type from to)
   (cond
    ((or (primitive-type? type)
@@ -2314,7 +2356,7 @@
        (values (reverse texps) subs pred))
       ((exp . rest)
        (define-values (texp subs* pred*)
-         (infer-exp exp env))
+         (infer exp env))
        (define-values (new-pred combined-subs)
          (eval-predicate* (predicate:compose pred pred*)
                           (compose-substitutions subs subs*)))
@@ -2326,7 +2368,7 @@
 (define (infer:if predicate consequent alternate env)
   ;; Infer predicate types and unify it with the boolean type.
   (define-values (predicate-texp predicate-subs predicate-pred)
-    (infer-exp predicate env))
+    (infer predicate env))
   (define predicate-unify-subs
     (unify (texp-types predicate-texp) (list type:bool)))
   ;; Combine the substitutions and apply them to the environment.
@@ -2337,13 +2379,13 @@
   ;; Infer consequent and alternate types and unify them against each
   ;; other.  Each branch of an 'if' should have the same type.
   (define-values (consequent-texp consequent-subs consequent-pred)
-    (infer-exp consequent env0))
+    (infer consequent env0))
   (define combined-subs-1
     (compose-substitutions combined-subs-0 consequent-subs))
   (define env1
     (apply-substitutions-to-env env0 consequent-subs))
   (define-values (alternate-texp alternate-subs alternate-pred)
-    (infer-exp alternate env1))
+    (infer alternate env1))
   (define combined-subs-2
     (compose-substitutions combined-subs-1 alternate-subs))
   ;; Eval combined predicate.
@@ -2371,7 +2413,7 @@
             (extend-env param type env*))
           env params param-types))
   (define-values (body* body-subs body-pred)
-    (infer-exp body env*))
+    (infer body env*))
   (define-values (pred subs)
     (eval-predicate* body-pred body-subs))
   (values (texp (list (generalize
@@ -2437,7 +2479,7 @@
   ;; The type signature of primitive functions can be looked up
   ;; directly in the environment.
   (define-values (operator* operator-subs operator-pred)
-    (infer-exp operator env))
+    (infer operator env))
   (define env*
     (apply-substitutions-to-env env operator-subs))
   ;; Infer the arguments.
@@ -2472,7 +2514,7 @@
 
 (define (infer:struct-ref exp field env)
   (define-values (exp* exp-subs exp-pred)
-    (infer-exp exp env))
+    (infer exp env))
   (define exp-type (single-type exp*))
   (define tvar (fresh-variable-type))
   (define-values (pred combined-subs)
@@ -2487,11 +2529,11 @@
 
 (define (infer:array-ref array-exp index-exp env)
   (define-values (array-exp* array-exp-subs array-exp-pred)
-    (infer-exp array-exp env))
+    (infer array-exp env))
   (define array-type (single-type array-exp*))
   (define env* (apply-substitutions-to-env env array-exp-subs))
   (define-values (index-exp* index-exp-subs index-exp-pred)
-    (infer-exp index-exp env*))
+    (infer index-exp env*))
   (define index-type (single-type index-exp*))
   (define combined-subs
     (compose-substitutions array-exp-subs index-exp-subs))
@@ -2523,7 +2565,7 @@
           names
           exp-types))
   (define-values (body* body-subs body-pred)
-    (infer-exp body env*))
+    (infer body env*))
   (define-values (pred combined-subs)
     (eval-predicate* (predicate:compose exp-pred body-pred)
                      (compose-substitutions exp-subs body-subs)))
@@ -2554,7 +2596,7 @@
           names
           exp-types))
   (define-values (body* body-subs body-pred)
-    (infer-exp body env*))
+    (infer body env*))
   (define-values (pred combined-subs)
     (eval-predicate* (predicate:compose exp-pred body-pred)
                      (compose-substitutions exp-subs body-subs)))
@@ -2613,7 +2655,7 @@
        (values (reverse texps) subs pred env))
       ((('function name exp) . rest)
        (define-values (texp subs* pred*)
-         (infer-exp exp env))
+         (infer exp env))
        (define-values (new-pred combined-subs)
          (eval-predicate* (predicate:compose pred pred*)
                           (compose-substitutions subs subs*)))
@@ -2646,7 +2688,7 @@
   (define-values (exps exp-subs exp-pred env*)
     (infer-bindings bindings '() '() predicate:succeed env))
   (define-values (body* body-subs body-pred)
-    (infer-exp body env*))
+    (infer body env*))
   (define-values (pred combined-subs)
     (eval-predicate* (predicate:compose exp-pred body-pred)
                      (compose-substitutions exp-subs body-subs)))
@@ -2666,7 +2708,7 @@
 ;; - a typed expression
 ;; - a list of substitutions
 ;; - a type predicate
-(define (infer-exp exp env)
+(define (infer exp env)
   (match exp
     ((? constant?)
      (infer:constant exp))
@@ -2698,40 +2740,21 @@
     ;; earlier compiler pass.
     (_ (error "unknown form" exp))))
 
-(define (type-descriptor->type desc)
-  (match desc
-    ('bool type:bool)
-    ('int type:int)
-    ('float type:float)
-    ('vec2 type:vec2)
-    ('vec3 type:vec3)
-    ('vec4 type:vec4)
-    ('mat3 type:mat3)
-    ('mat4 type:mat4)
-    ('sampler-2d type:sampler-2d)
-    (('array desc* (? exact-integer? length) (? exact-integer? rest) ...)
-     (let loop ((rest rest)
-                (prev (array-type (type-descriptor->type desc*) length)))
-       (match rest
-         (() prev)
-         ((length . rest)
-          (loop rest (array-type prev length))))))))
-
-(define (top-level-type-env stage)
-  (case stage
-    ((vertex)
-     `((vertex:position . ,type:vec4)
-       (vertex:point-size . ,type:float)
-       (vertex:clip-distance . ,type:float)))
-    ((fragment)
-     `((fragment:depth . ,type:float)
-       (fragment:coord . ,type:vec4)))))
+(define (infer:top-level-env stage)
+  (fold (lambda (v env)
+          (let ((name (seagull-variable-name v))
+                (type (seagull-variable-type v)))
+            (extend-env name type env)))
+        (empty-env)
+        (find-variables
+         (lambda (v)
+           (variable-for-stage? v stage)))))
 
 ;; TODO: Add some kind of context object that is threaded through the
 ;; inference process so that when a type error occurs we can show the
 ;; expression that caused it.
-(define (infer-types exp stage)
-  (infer-exp exp (top-level-type-env stage)))
+(define (infer* exp stage)
+  (infer exp (infer:top-level-env stage)))
 
 
 ;;;
@@ -2838,9 +2861,9 @@
                                            globals))
                            (match func
                              (('lambda _ body)
-                              (infer-exp (untype body)
-                                         (compose-envs env
-                                                       (top-level-type-env stage)))))
+                              (let ((top (infer:top-level-env stage)))
+                                (infer (untype body)
+                                       (compose-envs env top)))))
                            (define subs*
                              (compose-substitutions subs (vars->subs func env)))
                            (define func*
@@ -2870,28 +2893,26 @@
 (define (type-descriptor->glsl desc)
   (match desc
     ((? symbol?)
-     (match (type-descriptor->type desc)
+     (match (lookup-type desc)
        ((? primitive-type? primitive)
-        (primitive-type-name primitive))
+        (primitive-type-glsl-name primitive))
        ((? struct-type? struct)
-        (struct-type-name struct))))
+        (struct-type-glsl-name struct))))
     (('array desc* length)
      (format #f "~a[~a]"
              (type-descriptor->glsl desc*)
              length))))
 
-(define (type->type-descriptor type)
+(define (type->glsl type)
   (cond
    ((primitive-type? type)
-    (primitive-type-name type))
+    (primitive-type-glsl-name type))
    ((struct-type? type)
-    (struct-type-name type))
+    (struct-type-glsl-name type))
    ((array-type? type)
-    `(array ,(type->type-descriptor (array-type-ref type))
-            ,(array-type-length type)))))
-
-(define (type->glsl type)
-  (type-descriptor->glsl (type->type-descriptor type)))
+    (format #f "~a[~a]"
+            (type->glsl (array-type-ref type))
+            length))))
 
 (define (single-temp temps)
   (match temps
@@ -3226,7 +3247,7 @@
       (let* ((simplified (simplify expanded (empty-env)))
              (pruned (prune simplified))
              (hoisted (hoist-functions* pruned))
-             (inferred (infer-types hoisted stage))
+             (inferred (infer* hoisted stage))
              (resolved (resolve-overloads inferred stage)))
         (values resolved global-map (unique-identifier-counter))))))
 
@@ -3511,11 +3532,11 @@ Run the partial evaluator on EXP for shader STAGE."
 Run type inference on EXP for shader STAGE."
   (parameterize ((unique-identifier-counter 0))
     (pretty-print
-     (infer-types (hoist-functions*
-                   (prune
-                    (simplify*
-                     (expand* exp stage))))
-                  stage))))
+     (infer* (hoist-functions*
+              (prune
+               (simplify*
+                (expand* exp stage))))
+             stage))))
 
 (define-meta-command ((seagull-inspect chickadee) repl module)
   "seagull-inspect MODULE