path: root/chickadee/graphics/seagull.scm

;;; Chickadee Game Toolkit
;;; Copyright © 2023 David Thompson <davet@gnu.org>
;;;
;;; Chickadee is free software: you can redistribute it and/or modify
;;; it under the terms of the GNU General Public License as published
;;; by the Free Software Foundation, either version 3 of the License,
;;; or (at your option) any later version.
;;;
;;; Chickadee is distributed in the hope that it will be useful, but
;;; WITHOUT ANY WARRANTY; without even the implied warranty of
;;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
;;; General Public License for more details.
;;;
;;; You should have received a copy of the GNU General Public License
;;; along with this program.  If not, see
;;; <http://www.gnu.org/licenses/>.

;;; Commentary:
;;
;; The Seagull shading language is a purely functional, statically
;; typed, Scheme-like language that compiles to GLSL.  The reality of
;; how GPUs work imposes some significant language restrictions, but
;; they are restrictions anyone who writes shader code is already used
;; to.
;;
;; Features:
;; - Purely functional
;; - Statically typed via type inference
;; - Lexical scoping
;; - Nested functions
;; - Multiple return values
;; - Vertex and fragment shader output
;; - Emits code for multiple GLSL versions
;;
;; Limitations:
;; - First-order functions
;; - No closures
;; - No recursion
;;
;; TODO:
;; - Seagull unquote
;; - User defined structs
;; - Loops
;; - Better error messages (especially around type predicate failure)
;; - Helper function modules
;; - Shader composition
;; - Interpreter
;;
;;; Code:
(define-module (chickadee graphics seagull)
  #:use-module (chickadee graphics engine)
  #:use-module (chickadee graphics shader)
  #:use-module (ice-9 exceptions)
  #:use-module (ice-9 format)
  #:use-module (ice-9 match)
  #:use-module (ice-9 pretty-print)
  #:use-module ((rnrs base) #:select (mod))
  #:use-module (srfi srfi-1)
  #:use-module (srfi srfi-9)
  #:use-module (srfi srfi-9 gnu)
  #:use-module (srfi srfi-11)
  #:use-module (system repl command)
  #:use-module (system repl common)
  #:export (seagull-module?
            seagull-module-vertex?
            seagull-module-fragment?
            seagull-module-stage
            seagull-module-inputs
            seagull-module-outputs
            seagull-module-uniforms
            seagull-module-source
            seagull-module-compiled
            seagull-module-global-map
            compile-seagull-module
            compile-shader
            link-seagull-modules
            define-vertex-shader
            define-fragment-shader))

;; The Seagull compiler is designed as a series of source-to-source
;; program transformations (as described in "Compilation by Program
;; Transformation" by Richard Kelsey) in which each transformation
;; pass results in a program that is one step closer to being directly
;; emitted to GLSL code.
;;
;; I wouldn't have been able to write this compiler without the
;; fantastic "An Incremental Approach to Compiler Construction" paper
;; by Abdulaziz Ghuloum that showed me that even a mere mortal could
;; write a useful compiler.  Thanks to Christine Lemmer Webber for
;; pointing me to that paper.
;;
;; The first pass of the compiler, the expander, converts Seagull code
;; into an intermediate form that uses fewer syntactic forms and
;; renames all variables to be program-unique.  A simplifier pass then
;; propagates constants and evaluates expressions that can be computed
;; at compile-time such as (+ 1 2).  Each expression is then
;; associated with a type via a type inference pass.  The fully typed
;; program can then be emitted directly to GLSL code.  There are other
;; passes, but these are the most important.


;;;
;;; Compiler helpers
;;;

;; This is where we keep miscellaneous code that is useful for many
;; stages of the compiler.

(define (float? x)
  (and (number? x) (inexact? x)))

;; Constant types are fundamental data types that need no compilation.
(define (constant? x)
  (or (exact-integer? x)
      (float? x)
      (boolean? x)))

(define (top-level-qualifier? x)
  (memq x '(in out uniform)))


;;;
;;; Lexical environments
;;;

;; Environments keep track of the variables that are in scope of an
;; expression.

(define (empty-env)
  '())

(define-syntax-rule (new-env (key value) ...)
  (list (cons key value) ...))

(define &seagull-unbound-variable-error
  (make-exception-type '&seagull-unbound-variable-error &error '(name)))

(define make-seagull-unbound-variable-error
  (record-constructor &seagull-unbound-variable-error))

(define seagull-unbound-variable-name
  (exception-accessor &seagull-unbound-variable-error
                      (record-accessor &seagull-unbound-variable-error 'name)))

(define (lookup name env)
  (or (assq-ref env name)
      (raise-exception
       (make-exception
        (make-seagull-unbound-variable-error name)
        (make-exception-with-origin lookup)
        (make-exception-with-message "seagull: unbound variable")
        (make-exception-with-irritants (list name (env-names env)))))))

(define* (lookup* name env #:optional default)
  (let loop ((env env))
    (match env
      (() default)
      (((k . v) . rest)
       (if (eq? k name) v (loop rest))))))

(define (lookup-all names env)
  (map (lambda (name) (lookup name env)) names))

(define (extend-env name value env)
  (alist-cons name value env))

(define (compose-envs . envs)
  (concatenate envs))

(define (env-names env)
  (map car env))

(define (env-values env)
  (map cdr env))

(define (env-map proc env)
  (map (match-lambda
         ((name . exp)
          (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)
               (proc name exp)))
            env))


;;;
;;; Types
;;;

;; Record types are not used here because these type objects appear in
;; the compiled intermediate form of Seagull, which is generated at
;; compile-time.  Record types cannot be interned so simple tagged
;; lists are used instead.
(define-syntax-rule (define-symbolic-type name
                      constructor predicate (field getter) ...)
  (begin
    (define (constructor field ...)
      (list 'name field ...))
    (define (predicate obj)
      (match obj
        (('name field ...) #t)
        (_ #f)))
    (define (getter obj)
      (match obj
        (('name field ...)
         field)))
    ...))

(define-symbolic-type primitive
  primitive-type
  primitive-type?
  (name primitive-type-name)
  (glsl-name primitive-type-glsl-name))

(define-symbolic-type outputs
  outputs-type
  outputs-type?)

(define-symbolic-type struct
  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-symbolic-type array
  array-type
  array-type?
  (type array-type-ref)
  (length array-type-length))

(define unique-variable-type-counter (make-parameter 0))

(define (unique-variable-type-number)
  (let ((n (unique-variable-type-counter)))
    (unique-variable-type-counter (+ n 1))
    n))

(define (unique-variable-type-name)
  (string->symbol
   (format #f "T~a" (unique-variable-type-number))))

(define-symbolic-type tvar
  variable-type
  variable-type?
  (name variable-type-name))

(define (fresh-variable-type)
  (variable-type (unique-variable-type-name)))

(define (fresh-variable-types-for-list lst)
  (map (lambda (_x) (fresh-variable-type)) lst))

(define-symbolic-type ->
  function-type
  function-type?
  (parameters function-type-parameters)
  (returns function-type-returns))

;; For GLSL primitives that support multiple arities.
(define-symbolic-type case->
  function-case-type
  function-case-type?
  (cases function-case-type-cases))

(define (function-case-type-ref type arity)
  (assv-ref (function-case-type-cases type) arity))

;; For parametric polymorphism.
(define-symbolic-type for-all
  type-scheme
  type-scheme?
  (quantifiers type-scheme-quantifiers)
  (type type-scheme-ref))

;; For ad-hoc polymorphism.
(define-symbolic-type qualified
  qualified-type
  qualified-type?
  (type qualified-type-ref)
  (predicate qualified-type-predicate))

(define (type? obj)
  (or (primitive-type? obj)
      (variable-type? obj)
      (function-type? obj)
      (function-case-type? obj)
      (struct-type? obj)
      (outputs-type? obj)))

;; Type predicates represent additional constraints associated with a
;; typed expression.  They are used to implement ad-hoc polymorphism.
;; For example, a function with signature (-> (a a) (a)) could specify
;; the following predicate to allow either ints or floats as arguments:
;;
;;   (let ((a (variable-type 'T0)))
;;     (predicate:or (predicate:= a type:int)
;;                   (predicate:= a type:float)
(define-record-type <type-predicate>
  (make-type-predicate exp evaluator substituter)
  type-predicate?
  (exp type-predicate-exp) ; symbolic representation
  (evaluator type-predicate-evaluator) ; eval procedure
  (substituter type-predicate-substituter)) ; substitute procedure

(define (print-type-predicate pred port)
  (format port "#<type-predicate ~a>" (type-predicate-exp pred)))

(set-record-type-printer! <type-predicate> print-type-predicate)

(define predicate:succeed
  (make-type-predicate
   #t
   (lambda () (values predicate:succeed '()))
   (lambda (from to) predicate:succeed)))

(define (predicate:succeed? pred)
  (eq? pred predicate:succeed))

(define predicate:fail
  (make-type-predicate
   #f
   (lambda () (values predicate:fail '()))
   (lambda (from to) predicate:fail)))

(define (predicate:fail? pred)
  (eq? pred predicate:fail))

(define (predicate:= a b)
  (cond
   ((or (variable-type? a)
        (variable-type? b))
    (make-type-predicate
     `(= ,a ,b)
     (lambda ()
       (values (predicate:= a b) '()))
     (lambda (from to)
       (predicate:= (apply-substitution-to-type a from to)
                    (apply-substitution-to-type b from to)))))
   ((equal? a b)
    predicate:succeed)
   (else
    predicate:fail)))

(define (predicate:substitute a b)
  (make-type-predicate
   `(substitute ,a ,b)
   (lambda ()
     (values predicate:succeed (list (cons a b))))
   (lambda (from to)
     (predicate:substitute (apply-substitution-to-type a from to)
                           (apply-substitution-to-type b from to)))))

(define (predicate:struct-field struct field field-var)
  (cond
   ((struct-type? struct)
    (let ((field-type (struct-type-ref struct field)))
      (if field-type
          (predicate:substitute field-var field-type)
          predicate:fail)))
   ((variable-type? struct)
    (make-type-predicate
     `(struct-field ,struct ,field ,field-var)
     (lambda ()
       (values (predicate:struct-field struct field field-var) '()))
     (lambda (from to)
       (predicate:struct-field
        (apply-substitution-to-type struct from to)
        field
        (apply-substitution-to-type field-var from to)))))
   (else predicate:fail)))

(define (predicate:array-element array element-var)
  (cond
   ((array-type? array)
    (predicate:substitute element-var (array-type-ref array)))
   ((variable-type? array)
    (make-type-predicate
     `(array-element ,array ,element-var)
     (lambda ()
       (values (predicate:array-element array element-var) '()))
     (lambda (from to)
       (predicate:array-element
        (apply-substitution-to-type array from to)
        (apply-substitution-to-type element-var from to)))))
   (else predicate:fail)))

;; All the predicates must succeed, but unlike 'predicate:and', they
;; can succeed independently of one another.
(define (predicate:compose pred . rest)
  (if (null? rest)
      pred
      (let ((other (apply predicate:compose rest)))
        (cond
         ((and (predicate:succeed? pred) (predicate:succeed? other))
          predicate:succeed)
         ((predicate:succeed? pred)
          other)
         ((predicate:succeed? other)
          pred)
         (else
          (make-type-predicate
           `(compose ,(type-predicate-exp pred)
                     ,(type-predicate-exp other))
           (lambda ()
             (let-values (((pred* pred-subs) (eval-predicate pred)))
               (cond
                ;; Left succeeds, now check the right.
                ((predicate:succeed? pred*)
                 (let-values (((other* other-subs) (eval-predicate other)))
                   (cond
                    ;; Right succeeds, the composition is a success.
                    ((predicate:succeed? other*)
                     (values predicate:succeed
                             (compose-substitutions pred-subs other-subs)))
                    ;; Right fails, so the composition fails.
                    ((predicate:fail? other*)
                     (values predicate:fail '()))
                    ;; Also inconclusive, return the same composition.
                    (else
                     (values other pred-subs)))))
                ;; Left fails, so the composition fails.
                ((predicate:fail? pred*)
                 (values predicate:fail '()))
                ;; Left predicate is inconclusive, try the right.
                (else
                 (let-values (((other* other-subs) (eval-predicate other)))
                   (cond
                    ;; Right succeeds, return the left.
                    ((predicate:succeed? other*)
                     (values pred other-subs))
                    ;; Right fails, so the composition fails.
                    ((predicate:fail? other*)
                     (values predicate:fail '()))
                    ;; Also inconclusive, return the same composition.
                    (else
                     (values (predicate:compose pred other) '()))))))))
           (lambda (from to)
             (predicate:compose
              (apply-substitution-to-predicate pred from to)
              (apply-substitution-to-predicate other from to)))))))))

;; The 'and' predicate succeeds if and when all the given predicates
;; succeed.
(define (predicate:and . preds)
  (match (remove predicate:succeed? preds)
    (()
     predicate:succeed)
    ((pred)
     pred)
    ((or ((? predicate:fail?) _)
         (_ (? predicate:fail?)))
     predicate:fail)
    ((a b)
     (make-type-predicate
      `(and ,(type-predicate-exp a) ,(type-predicate-exp b))
      (lambda ()
        (let-values (((a* a-subs) (eval-predicate a)))
          (cond
           ;; Left succeeds, now try the right.
           ((predicate:succeed? a*)
            (let-values (((b* b-subs) (eval-predicate b)))
              (cond
               ;; Right succeeds, so the 'and' succeeds.
               ((predicate:succeed? b*)
                (values predicate:succeed
                        (compose-substitutions a-subs b-subs)))
               ;; Right fails, so the 'and' fails.
               ((predicate:fail? b*)
                (values predicate:fail '()))
               ;; Right is inconclusive, so return the same 'and'.
               (else
                (predicate:and a b)))))
           ;; Left fails, so the 'and' fails.
           ((predicate:fail? a*)
            (values predicate:fail '()))
           ;; Left is inconclusive, so return the same 'and'.
           (else
            (values (predicate:and a b) '())))))
      (lambda (from to)
        (predicate:and (apply-substitution-to-predicate a from to)
                       (apply-substitution-to-predicate b from to)))))
    ((a . rest)
     (predicate:and a (apply predicate:and rest)))))

;; The 'or' predicate succeeds if and when any given predicate
;; succeeds.
(define (predicate:or . preds)
  (match (remove predicate:fail? preds)
    (()
     predicate:fail)
    ((pred)
     pred)
    ((or ((? predicate:succeed?) _)
         (_ (? predicate:succeed?)))
     predicate:succeed)
    ((a b)
     (make-type-predicate
      `(or ,(type-predicate-exp a) ,(type-predicate-exp b))
      (lambda ()
        (let-values (((a* a-subs) (eval-predicate a)))
          (cond
           ;; Left succeeds, so the 'or' succeeds.
           ((predicate:succeed? a*)
            (values predicate:succeed a-subs))
           ;; Left fails, so remove the 'or' and eval the right.
           ((predicate:fail? a*)
            (eval-predicate b))
           ;; Left is inconclusive, check the right.
           (else
            (let-values (((b* b-subs) (eval-predicate b)))
              (cond
               ;; Right succeeds, so the 'or' succeeds.
               ((predicate:succeed? b*)
                (values predicate:succeed b-subs))
               ;; Right fails, so remove the 'or' and return the left.
               ((predicate:fail? b*)
                (values a '()))
               (else
                (values (predicate:or a b) '()))))))))
      (lambda (from to)
        (predicate:or (apply-substitution-to-predicate a from to)
                      (apply-substitution-to-predicate b from to)))))
    ((a . rest)
     (predicate:or a (apply predicate:or rest)))))

(define (predicate:any var . types)
  (apply predicate:or
         (map (lambda (type)
                (predicate:= var type))
              types)))

(define (apply-substitution-to-predicate pred from to)
  ((type-predicate-substituter pred) from to))

(define (apply-substitutions-to-predicate pred subs)
  (env-fold (lambda (from to pred*)
              (apply-substitution-to-predicate pred* from to))
            pred
            subs))

(define (eval-predicate pred)
  ((type-predicate-evaluator pred)))

(define (eval-predicate* pred subs)
  (define-values (new-pred pred-subs)
    (eval-predicate
     (apply-substitutions-to-predicate pred subs)))
  ;; TODO: Get information about *why* the predicate failed.
  (unless new-pred
    (seagull-type-error "type predicate failed" (list pred) eval-predicate*))
  ;; Recursively evaluate the predicate, applying the substitutions
  ;; generated by the last evaluation, until it cannot be simplified
  ;; any further.
  (if (null? pred-subs)
      (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-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))
(define-struct-type (type:vec3 vec3)
  (type:float x)
  (type:float y)
  (type:float z))
(define-struct-type (type:vec4 vec4)
  (type:float x)
  (type:float y)
  (type:float z)
  (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-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))


;;;
;;; Built-in variables
;;;

(define-record-type <seagull-variable>
  (%make-seagull-variable name glsl-name type stages qualifier)
  seagull-variable?
  (name seagull-variable-name)
  (glsl-name seagull-variable-glsl-name)
  (type seagull-variable-type)
  (stages seagull-variable-stages)
  (qualifier seagull-variable-qualifier))

(define* (make-seagull-variable name #:key glsl-name type stages qualifier)
  (%make-seagull-variable name glsl-name type stages qualifier))

(define (output-variable? variable)
  (eq? (seagull-variable-qualifier variable) 'output))

(define (input-variable? variable)
  (eq? (seagull-variable-qualifier variable) 'input))

(define (variable-for-stage? variable stage)
  (memq stage (seagull-variable-stages variable)))

(define *seagull-variables* (make-hash-table))

(define (register-seagull-variable! variable)
  (hashq-set! *seagull-variables*
              (seagull-variable-name variable)
              variable))

(define (find-variables pred)
  (hash-fold (lambda (k v memo)
               (if (pred v)
                   (cons v memo)
                   memo))
             '()
             *seagull-variables*))

(define (lookup-output-variable name)
  (let ((variable (hashq-ref *seagull-variables* name)))
    (and (seagull-variable? variable)
         (output-variable? variable)
         variable)))

(define (lookup-output-variable-for-stage name stage)
  (let ((variable (lookup-output-variable name)))
    (and (seagull-variable? variable)
         (variable-for-stage? variable stage)
         variable)))

(define (lookup-input-variable name)
  (let ((variable (hashq-ref *seagull-variables* name)))
    (and (seagull-variable? variable)
         (input-variable? variable)
         variable)))

(define-syntax-rule (define-seagull-variable name args ...)
  (register-seagull-variable! (make-seagull-variable 'name args ...)))

(define-seagull-variable vertex:position
  #:glsl-name "gl_Position"
  #:type type:vec4
  #:stages '(vertex)
  #:qualifier 'output)

(define-seagull-variable vertex:point-size
  #:glsl-name "gl_PointSize"
  #:type type:float
  #:stages '(vertex)
  #:qualifier 'output)

(define-seagull-variable vertex:clip-distance
  #:glsl-name "gl_ClipDistance"
  #:type type:float
  #:stages '(vertex)
  #:qualifier 'output)

(define-seagull-variable fragment:depth
  #:glsl-name "gl_FragDepth"
  #:type type:float
  #:stages '(fragment)
  #:qualifier 'output)

(define-seagull-variable fragment:coord
  #:glsl-name "gl_FragCoord"
  #:type type:vec4
  #:stages '(fragment)
  #:qualifier 'input)


;;;
;;; Primitives
;;;

(define-record-type <seagull-primitive>
  (%make-seagull-primitive name glsl-name stages type proc expand emit)
  seagull-primitive?
  (name seagull-primitive-name)
  (glsl-name seagull-primitive-glsl-name)
  (stages seagull-primitive-stages)
  (type seagull-primitive-type)
  (proc seagull-primitive-proc)
  (expand seagull-primitive-expand)
  (emit seagull-primitive-emit))

(define (make-default-expander name)
  (define (expand:default args stage env)
    `(primcall ,name ,@(expand:list args stage env)))
  expand:default)

(define (make-default-emitter name)
  (define (emit:default args port)
    (format port "~a(~a)"
            name
            (string-join (map symbol->string args) ", ")))
  emit:default)

(define (make-infix-emitter name)
  (define (emit:infix args port)
    (match args
      ((a b)
       (format port "~a ~a ~a" a name b))))
  emit:infix)

(define* (make-seagull-primitive #:key name type proc
                                 (glsl-name name)
                                 (stages '(vertex fragment))
                                 (expand (make-default-expander name))
                                 (emit (make-default-emitter glsl-name)))
  (%make-seagull-primitive name glsl-name stages type proc expand emit))

(define *seagull-primitives* (make-hash-table))

(define (register-seagull-primitive! primitive)
  (hashq-set! *seagull-primitives*
              (seagull-primitive-name primitive)
              primitive)
  *unspecified*)

(define (lookup-seagull-primitive name)
  (hashq-ref *seagull-primitives* name))

(define-syntax-rule (define-seagull-primitive name args ...)
  (register-seagull-primitive!
   (make-seagull-primitive #:name 'name args ...)))

(define (primitive-call? x stage)
  (let ((primitive (lookup-seagull-primitive x)))
    (and (seagull-primitive? primitive)
         (memq stage (seagull-primitive-stages primitive)))))

(define-syntax-rule (-> (params ...) (returns ...))
  (function-type (list params ...) (list returns ...)))

(define-syntax-rule (->case (arity type) ...)
  (function-case-type `((arity . ,type) ...)))

(define-syntax overload
  (syntax-rules (->)
    ((_ ((var types ...) ...) (-> (args ...) (returns ...)))
     (parameterize ((unique-variable-type-counter 0))
       (let ((var (fresh-variable-type)) ...)
         (type-scheme
          (list var ...)
          (qualified-type
           (function-type (list args ...) (list returns ...))
           (predicate:compose
            (predicate:any var types ...) ...))))))))

(define-syntax-rule (a+b->c (ta tb tc) ...)
  (parameterize ((unique-variable-type-counter 0))
    (let ((a (fresh-variable-type))
          (b (fresh-variable-type))
          (c (fresh-variable-type)))
      (type-scheme
       (list a b c)
       (qualified-type
        (function-type (list a b) (list c))
        (predicate:or
         (predicate:and (predicate:= a ta)
                        (predicate:= b tb)
                        (predicate:substitute c tc))
         ...))))))

(define-syntax-rule (a+b+c->d (ta tb tc td) ...)
  (parameterize ((unique-variable-type-counter 0))
    (let ((a (fresh-variable-type))
          (b (fresh-variable-type))
          (c (fresh-variable-type))
          (d (fresh-variable-type)))
      (type-scheme
       (list a b c d)
       (qualified-type
        (function-type (list a b c) (list d))
        (predicate:or
         (predicate:and (predicate:= a ta)
                        (predicate:= b tb)
                        (predicate:= c tc)
                        (predicate:substitute d td))
         ...))))))

(define-seagull-primitive +
  #:type (overload ((a type:int type:float
                       type:vec2 type:vec3 type:vec4
                       type:mat3 type:mat4))
                   (-> (a a) (a)))
  #:proc +
  #:expand
  (lambda (args stage env)
    (let loop ((args args))
      (match args
        (() 0)
        ((n) (expand n stage env))
        ((n . rest)
         `(primcall + ,(expand n stage env) ,(loop rest))))))
  #:emit (make-infix-emitter '+))

(define-seagull-primitive -
  #:type (overload ((a type:int type:float
                       type:vec2 type:vec3 type:vec4
                       type:mat3 type:mat4))
                   (-> (a a) (a)))
  #:proc -
  #:expand
  (lambda (args stage env)
    (let loop ((args args))
      (match args
        ((n) `(primcall - ,(expand n stage env) 0))
        ((m n)
         `(primcall - ,(expand m stage env) ,(expand n stage env)))
        ((n . rest)
         `(primcall - ,(expand n stage env) ,(loop rest))))))
  #:emit (make-infix-emitter '-))

(define-seagull-primitive *
  #:type (a+b->c (type:int type:int type:int)
                 (type:float type:float type:float)
                 (type:int type:float type:float)
                 (type:float type:int type:float)
                 (type:vec2 type:vec2 type:vec2)
                 (type:vec2 type:float type:vec2)
                 (type:float type:vec2 type:vec2)
                 (type:vec3 type:vec3 type:vec3)
                 (type:vec3 type:float type:vec3)
                 (type:float type:vec3 type:vec3)
                 (type:vec4 type:vec4 type:vec4)
                 (type:vec4 type:float type:vec4)
                 (type:float type:vec4 type:vec4)
                 (type:mat3 type:mat3 type:mat3)
                 (type:mat3 type:vec3 type:vec3)
                 (type:vec3 type:mat3 type:vec3)
                 (type:mat4 type:mat4 type:mat4)
                 (type:mat4 type:vec4 type:vec4)
                 (type:vec4 type:mat4 type:vec4))
  #:proc *
  #:expand
  (lambda (args stage env)
    (let loop ((args args))
      (match args
        (() 1)
        ((n) (expand n stage env))
        ((n . rest)
         `(primcall * ,(expand n stage env) ,(loop rest))))))
  #:emit (make-infix-emitter '*))

(define-seagull-primitive /
  #:type (a+b->c (type:int type:int type:int)
                 (type:float type:float type:float)
                 (type:float type:int type:float)
                 (type:int type:float type:float)
                 (type:vec2 type:vec2 type:vec2)
                 (type:vec2 type:float type:vec2)
                 (type:vec3 type:vec3 type:vec3)
                 (type:vec3 type:float type:vec3)
                 (type:vec4 type:vec4 type:vec4)
                 (type:vec4 type:float type:vec4)
                 (type:mat3 type:float type:mat3)
                 (type:mat4 type:float type:mat4))
  ;; The division of two integers can result in a rational,
  ;; non-integer, such as 1/2.  This isn't how integer division works
  ;; in GLSL, so we need to round the result to an integer.
  #:proc
  (lambda (x y)
    (let ((result (/ x y)))
      (if (or (float? result) (integer? result))
          result
          (round result))))
  #:expand
  (lambda (args stage env)
    (match args
      ((n)
       `(primcall / 1 ,(expand n stage env)))
      ((m n)
       `(primcall / ,(expand m stage env) ,(expand n stage env)))
      ((m n . rest)
       (let loop ((rest rest)
                  (exp `(primcall / ,(expand m stage env) ,(expand n stage env))))
         (match rest
           ((l)
            `(primcall / ,exp ,(expand l stage env)))
           ((l . rest)
            (loop rest `(primcall / ,exp ,(expand l stage env)))))))))
  #:emit (make-infix-emitter '/))

(define-seagull-primitive mod
  #:type (a+b->c (type:float type:float type:float)
                 (type:int type:int type:float)
                 (type:vec2 type:vec2 type:vec2)
                 (type:vec3 type:vec3 type:vec3)
                 (type:vec4 type:vec4 type:vec4)
                 (type:vec2 type:float type:vec2)
                 (type:vec3 type:float type:vec3)
                 (type:vec4 type:float type:vec4))
  #:proc mod)

(define-seagull-primitive floor
  #:type (overload ((a type:float type:vec2 type:vec3 type:vec4))
                   (-> (a) (a)))
  #:proc floor)

(define-seagull-primitive ceiling
  #:glsl-name 'ceil
  #:type (overload ((a type:float type:vec2 type:vec3 type:vec4))
                   (-> (a) (a)))
  #:proc ceiling)

(define-seagull-primitive int->float
  #:glsl-name 'float
  #:type (-> (type:int) (type:float))
  #:proc exact->inexact)

(define-seagull-primitive float->int
  #:glsl-name 'int
  #:type (-> (type:float) (type:int))
  #:proc (compose inexact->exact floor))

(define (make-comparison-expander name)
  (lambda (args stage env)
    (match args
      (() #t)
      ((x)
       (expand `(let ((x* ,x)) (,name x* x*)) stage env))
      ((x y)
       `(primcall ,name ,(expand x stage env) ,(expand y stage env)))
      ((x y . rest)
       (expand `(let ((y* ,y))
                  (and (,name ,x y*)
                       ,(let loop ((rest rest)
                                   (prev 'y*))
                          (match rest
                            ((z)
                             (list name prev z))
                            ((z . rest)
                             `(let ((z* ,z))
                                (and (,name ,prev z*)
                                     ,(loop rest 'z*))))))))
               stage env)))))

(define-syntax define-comparison-primitive
  (syntax-rules ()
    ((_ name)
     (define-comparison-primitive name name
       (make-comparison-expander 'name)))
    ((_ name glsl-name)
     (define-comparison-primitive name glsl-name
       (make-comparison-expander 'name)))
    ((_ name glsl-name expand)
     (define-seagull-primitive name
       #:glsl-name 'glsl-name
       #:type
       (overload ((a type:int type:float))
                 (-> (a a) (type:bool)))
       #:proc name
       #:expand expand
       #:emit (make-infix-emitter 'glsl-name)))))

(define-comparison-primitive = ==
  (lambda (args stage env)
    (match args
      (() #t)
      ((x)
       (expand `(let ((x* ,x)) (= x* x*)) stage env))
      ((x y)
       `(primcall = ,(expand x stage env) ,(expand y stage env)))
      ((x . rest)
       (expand `(let ((x* ,x))
                  (and ,@(map (lambda (y) `(= x* ,y)) rest)))
               stage env)))))

(define-comparison-primitive <)
(define-comparison-primitive <=)
(define-comparison-primitive >)
(define-comparison-primitive >=)

(define-seagull-primitive not
  #:glsl-name '!
  #:type (-> (type:bool) (type:bool))
  #:emit
  (lambda (args port)
    (match args
      ((a)
       (format port "!(~a)" a)))))

(define-seagull-primitive vec2
  #:type (->case
          (1 (-> (type:float) (type:vec2)))
          (2 (-> (type:float type:float) (type:vec2)))))

(define-seagull-primitive vec3
  #:type (->case
          (1 (-> (type:float) (type:vec3)))
          (2 (a+b->c (type:float type:vec2 type:vec3)
                     (type:vec2 type:float type:vec3)))
          (3 (-> (type:float type:float type:float) (type:vec3)))))

(define-seagull-primitive vec4
  #:type (->case
          (1 (-> (type:float) (type:vec4)))
          (2 (a+b->c (type:vec2 type:vec2 type:vec4)
                     (type:vec3 type:float type:vec4)
                     (type:float type:vec3 type:vec4)))
          (3 (a+b+c->d (type:vec2 type:float type:float type:vec4)
                       (type:float type:vec2 type:float type:vec4)
                       (type:float type:float type:vec2 type:vec4)))
          (4 (-> (type:float type:float type:float type:float) (type:vec4)))))

(define-seagull-primitive length
  #:type (overload ((a type:float type:vec2 type:vec3 type:vec4))
                   (-> (a) (type:float))))

(define-seagull-primitive abs
  #:type (overload ((a type:int type:float))
                   (-> (a) (a)))
  #:proc abs)

(define-seagull-primitive sqrt
  #:type (-> (type:float) (type:float))
  #:proc sqrt)

(define-seagull-primitive expt
  #:glsl-name 'pow
  #:type (overload ((a type:float type:vec2 type:vec3 type:vec4))
                   (-> (a a) (a)))
  #:proc expt)

(define-seagull-primitive min
  #:type (overload ((a type:int type:float)) (-> (a a) (a)))
  #:proc min)

(define-seagull-primitive max
  #:type (overload ((a type:int type:float)) (-> (a a) (a)))
  #:proc min)

(define-seagull-primitive sin
  #:type (-> (type:float) (type:float))
  #:proc sin)

(define-seagull-primitive cos
  #:type (-> (type:float) (type:float))
  #:proc cos)

(define-seagull-primitive tan
  #:type (-> (type:float) (type:float))
  #:proc tan)

(define-seagull-primitive clamp
  #:type (overload ((a type:int type:float)) (-> (a a a) (a))))

(define-seagull-primitive mix
  #:type (overload ((a type:int type:float type:vec2 type:vec3 type:vec4))
                   (-> (a a type:float) (a))))

(define-seagull-primitive step
  #:type (a+b->c (type:float type:float type:float)
                 (type:vec2 type:vec2 type:vec2)
                 (type:vec3 type:vec3 type:vec3)
                 (type:vec4 type:vec4 type:vec4)
                 (type:float type:vec2 type:vec2)
                 (type:float type:vec3 type:vec3)
                 (type:float type:vec4 type:vec4)))

(define-seagull-primitive smoothstep
  #:type (a+b+c->d (type:float type:float type:float type:float)
                   (type:vec2 type:vec2 type:vec2 type:vec2)
                   (type:vec3 type:vec3 type:vec3 type:vec3)
                   (type:vec4 type:vec4 type:vec4 type:vec4)
                   (type:float type:float type:vec2 type:vec2)
                   (type:float type:float type:vec3 type:vec3)
                   (type:float type:float type:vec4 type:vec4)))

(define-seagull-primitive texture
  #:stages '(fragment)
  #:type (-> (type:sampler-2d type:vec2) (type:vec4)))


;;;
;;; Macro expansion and alpha conversion
;;;

;; Macro expansion converts convenient but non-primitive syntax forms
;; (such as let*) into primitive syntax.  Seagull does not currently
;; support user defined macros, just a set of built-ins.
;;
;; Alpha conversion is the process of converting all the user defined
;; identifiers in a program to uniquely named identifiers.  This
;; process frees the compiler from having to worry about things like
;; '+' being a user defined variable that shadows the primitive
;; addition operation.

(define &seagull-syntax-error
  (make-exception-type '&seagull-syntax-error &error '(form)))

(define make-seagull-syntax-error
  (record-constructor &seagull-syntax-error))

(define seagull-syntax-form
  (exception-accessor &seagull-syntax-error
                      (record-accessor &seagull-syntax-error 'form)))

(define (seagull-syntax-error msg exp origin)
  (raise-exception
   (make-exception
    (make-seagull-syntax-error exp)
    (make-exception-with-origin origin)
    (make-exception-with-message
     (format #f "seagull syntax error: ~a" msg))
    (make-exception-with-irritants (list exp)))))

(define unique-identifier-counter (make-parameter 0))

(define (unique-identifier-number)
  (let ((n (unique-identifier-counter)))
    (unique-identifier-counter (+ n 1))
    n))

(define (unique-identifier)
  (string->symbol
   (format #f "V~a" (unique-identifier-number))))

(define (unique-identifiers-for-list lst)
  (map (lambda (_x) (unique-identifier)) lst))

(define (expand:top-level-env stage)
  (fold (lambda (v env)
          (let ((name (seagull-variable-name v)))
            (extend-env name name env)))
        (empty-env)
        (find-variables
         (lambda (v)
           (and (input-variable? v)
                (variable-for-stage? v stage))))))

(define (alpha-convert names)
  (define names* (map (lambda (_name) (unique-identifier)) names))
  (fold extend-env (empty-env) names names*))

(define (expand:list exps stage env)
  (map (lambda (exp) (expand exp stage env)) exps))

(define (expand:variable exp stage env)
  (lookup exp env))

(define (expand:if predicate consequent alternate stage env)
  `(if ,(expand predicate stage env)
       ,(expand consequent stage env)
       ,(expand alternate stage env)))

(define (expand:let names exps body stage env)
  (if (null? names)
      (expand body stage env)
      (let* ((exps* (map (lambda (exp) (expand exp stage env)) exps))
             (env* (compose-envs (alpha-convert names) env))
             (bindings* (map list (lookup-all names env*) exps*)))
        `(let ,bindings* ,(expand `(begin ,@body) stage env*)))))

(define (expand:let-values names exps body stage env)
  (if (null? names)
      (expand body stage env)
      (let* ((exps* (map (lambda (exp)
                           (expand exp stage env))
                         exps))
             (env* (fold (lambda (names* env*)
                          (compose-envs (alpha-convert names*) env*))
                         env
                         names))
             (bindings* (map list
                             (map (lambda (names*)
                                    (lookup-all names* env*))
                                  names)
                             exps*)))
        `(let-values ,bindings* ,(expand `(begin ,@body) stage env*)))))

(define (expand:let* bindings body stage env)
  (match bindings
    (() (expand body stage env))
    ((binding . rest)
     (expand `(let (,binding)
                (let* ,rest ,body))
             stage
             env))))

(define (expand:let*-values bindings body stage env)
  (match bindings
    (() (expand body stage env))
    ((binding . rest)
     (expand `(let-values (,binding)
                (let*-values ,rest ,body))
             stage
             env))))

(define (expand:lambda params body stage env)
  (define env* (compose-envs (alpha-convert params) env))
  (define params* (lookup-all params env*))
  `(lambda ,params* ,(expand `(begin ,@body) stage env*)))

(define (expand:values exps stage env)
  `(values ,@(expand:list exps stage env)))

(define (expand:-> exp fields stage env)
  (define exp* (expand exp stage env))
  (match fields
    ((field . rest)
     (let loop ((fields rest)
                (prev `(struct-ref ,exp* ,field)))
       (match fields
         (() prev)
         ((next . rest)
          (loop `(struct-ref ,prev ,next)
                rest)))))))

(define (expand:@ exp indices stage env)
  (define exp* (expand exp stage env))
  (match indices
    ((i . rest)
     (let loop ((indices rest)
                (prev `(array-ref ,exp* ,(expand i stage env))))
       (match indices
         (() prev)
         ((j . rest)
          (loop `(array-ref ,prev ,(expand j stage env))
                rest)))))))

(define (expand:begin body stage env)
  (match body
    (((and ('define _ ...) definitions) ... body)
     (define bindings
       (map (match-lambda
              (('define (proc-name (? symbol? params) ...) body)
               (list proc-name (expand `(lambda ,params ,body) stage env)))
              (('define (? symbol? var-name) val)
               (list var-name (expand val stage env)))
              (invalid
               (seagull-syntax-error "invalid definition" invalid expand:begin)))
            definitions))
     (define names (map first bindings))
     (define env* (compose-envs (alpha-convert names) env))
     (let loop ((bindings bindings))
       (match bindings
         (()
          (expand body stage env*))
         (((name value) . rest)
          `(let ((,(lookup name env*) ,value)) ,(loop rest))))))
    (_
     (seagull-syntax-error "invalid begin form" `(begin ,@body) expand:begin))))

(define (expand:or exps stage env)
  (match exps
    (() #f)
    ((exp . rest)
     (expand `(let ((x ,exp)) (if x x (or ,@rest)))
             stage env))))

(define (expand:and exps stage env)
  (match exps
    (() #t)
    ((exp . rest)
     (expand `(let ((x ,exp)) (if x (and ,@rest) #f))
             stage env))))

(define (expand:cond clauses stage env)
  (define (cond->if clauses*)
    (match clauses*
      ;; Our version of 'cond' requires a final else clause because the
      ;; static type checker enforces that both branches of an 'if' must
      ;; be the same type.  If 'else' were optional then we wouldn't
      ;; know what type the final alternate branch should be.
      ((('else exp))
       exp)
      (((predicate consequent) . rest)
       `(if ,predicate
            ,consequent
            ,(cond->if rest)))
      (()
       (seagull-syntax-error "'cond' form must end with 'else' clause"
                             `(cond ,@clauses)
                             expand:cond))
      (_
       (seagull-syntax-error "invalid 'cond' form"
                             `(cond ,@clauses)
                             expand:cond))))
  (expand (cond->if clauses) stage env))

(define (expand:case key clauses stage env)
  (define (case->if clauses*)
    (match clauses*
      ;; Like 'cond', 'case' also requires a final 'else' clause.
      ((('else exp))
       exp)
      ((((possibilities ..1) consequent) . rest)
       `(if (or ,@(map (lambda (n) `(= key ,n)) possibilities))
            ,consequent
            ,(case->if rest)))
      (()
       (seagull-syntax-error "'case' form must end with 'else' clause"
                             `(case ,key ,@clauses)
                             expand:case))
      (_
       (seagull-syntax-error "invalid 'cond' form"
                             `(case ,key ,@clauses)
                             expand:case))))
  (expand `(let ((key ,key)) ,(case->if clauses)) stage env))

(define (expand:primitive-call operator operands stage env)
  (let ((primitive (lookup-seagull-primitive operator)))
    ((seagull-primitive-expand primitive) operands stage env)))

(define (expand:call operator operands stage env)
  `(call ,(expand operator stage env) ,@(expand:list operands stage env)))

(define (expand:top-level qualifiers types names body stage env)
  (let* ((global-map (alpha-convert names))
         (env* (compose-envs global-map env)))
    ;; TODO: Support interpolation qualifiers.
    (values `(top-level ,(map (lambda (qualifier type name)
                                (list qualifier type (lookup name env*)))
                              qualifiers types names)
                        ,(expand body stage env*))
            global-map)))

(define (expand:outputs names exps stage env)
  `(outputs
    ,@(map (lambda (name exp)
             (let ((variable (lookup-output-variable-for-stage name stage)))
               (list (if (seagull-variable? variable)
                         (seagull-variable-name variable)
                         ;; TODO: Check that the variable is an output
                         ;; variable.
                         (lookup name env))
                     (expand exp stage env))))
           names exps)))

(define (expand:discard stage env)
  (if (eq? stage 'fragment)
      '(outputs)
      (seagull-syntax-error "discard not allowed in vertex shader" exp expand)))

(define (expand exp stage env)
  (define (primitive-call-for-stage? x)
    (primitive-call? x stage))
  (match exp
    ;; Constants and variables:
    ((? constant?)
     exp)
    ((? symbol?)
     (expand:variable exp stage env))
    ;; Primitive syntax forms:
    (('if predicate consequent alternate)
     (expand:if predicate consequent alternate stage env))
    (('let (((? symbol? names) exps) ...) body ...)
     (expand:let names exps body stage env))
    (('let-values ((((? symbol? names) ...) exps) ...) body ...)
     (expand:let-values names exps body stage env))
    (('lambda ((? symbol? params) ...) body ...)
     (expand:lambda params body stage env))
    (('values exps ...)
     (expand:values exps stage env))
    (('outputs ((? symbol? names) exps) ...)
     (expand:outputs names exps stage env))
    (('top-level (((? top-level-qualifier? qualifiers) types names) ...)
                 body)
     (expand:top-level qualifiers types names body stage env))
    ;; Macros:
    (('-> exp (? symbol? members) ..1)
     (expand:-> exp members stage env))
    (('@ exp indices ...)
     (expand:@ exp indices stage env))
    (('begin body ...)
     (expand:begin body stage env))
    (('let* (bindings ...) body)
     (expand:let* bindings body stage env))
    (('let*-values (bindings ...) body)
     (expand:let*-values bindings body stage env))
    (('or exps ...)
     (expand:or exps stage env))
    (('and exps ...)
     (expand:and exps stage env))
    (('cond clauses ...)
     (expand:cond clauses stage env))
    (('case key clauses ...)
     (expand:case key clauses stage env))
    (('discard)
     (expand:discard stage env))
    ;; Primitive calls:
    (((? primitive-call-for-stage? operator) args ...)
     (expand:primitive-call operator args stage env))
    ;; Function calls:
    ((operator args ...)
     (expand:call operator args stage env))
    ;; Syntax error:
    (_
     (seagull-syntax-error "unknown form" exp expand))))

(define (expand* exp stage)
  (expand exp stage (expand:top-level-env stage)))


;;;
;;; Constant propagation and folding
;;;

;; Replace references to constants (variables that store an constant
;; value: integer, float, boolean) with the constants themselves.
;; Then look for opportunities to evaluate primitive expressions that
;; have constant arguments.  This will make the type inferencer's job
;; a bit easier.

(define (simplify:list exps env)
  (map (lambda (exp) (simplify exp env)) exps))

(define (simplify:if predicate consequent alternate env)
  (define predicate* (simplify predicate env))
  (cond
   ((eq? predicate* #t)
    (simplify consequent env))
   ((eq? predicate* #f)
    (simplify alternate env))
   (else
    `(if ,predicate*
         ,(simplify consequent env)
         ,(simplify alternate env)))))

(define (simplify:lambda params body env)
  `(lambda ,params ,(simplify body env)))

(define (simplify:values exps env)
  `(values ,@(simplify:list exps env)))

(define (simplify:let names exps body env)
  (define exps* (simplify:list exps env))
  ;; Extend environment with known constants.
  (define env*
    (fold (lambda (name exp env*)
            (if (constant? exp)
                (extend-env name exp env*)
                env*))
          env names exps*))
  ;; Drop all bindings for constant expressions.
  (define bindings
    (filter-map (lambda (name exp)
                  (if (constant? exp)
                      #f
                      (list name exp)))
                names exps*))
  ;; If there are no bindings left, remove the 'let' entirely.
  (if (null? bindings)
      (simplify body env*)
      `(let ,bindings
         ,(simplify body env*))))

(define (simplify:let-values names exps body env)
  (define exps* (simplify:list exps env))
  ;; Extend environment with known constants.
  (define env*
    (fold (lambda (names exp env)
            (match exp
              ((? constant?)
               (match names
                 ((name)
                  (extend-env name exp env))))
              (('values vals ...)
               (fold (lambda (name val env)
                       (if (constant? val)
                           (extend-env name val env)
                           env))
                     env names vals))
              (_ env)))
          env names exps*))
  ;; Drop all bindings for constant expressions.
  (define bindings
    (filter-map (lambda (names exp)
                  (match exp
                    ((? constant?) #f)
                    (('values vals ...)
                     (define-values (names* exps*)
                       (unzip2
                        (filter-map (lambda (name val)
                                      (if (constant? val)
                                          #f
                                          (list name val)))
                                    names vals)))
                     (if (null? names*)
                         #f
                         (list names* exps*)))
                    (_ (list names exp))))
                names exps*))
  ;; If there are no bindings left, remove the 'let' entirely.
  (if (null? bindings)
      (simplify body env*)
      `(let-values ,bindings
         ,(simplify body env*))))

(define (simplify:primcall op args env)
  (let* ((primitive (lookup-seagull-primitive op))
         (proc (seagull-primitive-proc primitive))
         (args* (simplify:list args env)))
    (if (and (procedure? proc) (every constant? args*))
        (apply proc args*)
        `(primcall ,op ,@args*))))

(define (simplify:call operator args env)
  `(call ,(simplify operator env)
         ,@(simplify:list args env)))

(define (simplify:struct-ref exp field env)
  `(struct-ref ,(simplify exp env) ,field))

(define (simplify:array-ref array-exp index-exp env)
  `(array-ref ,(simplify array-exp env)
              ,(simplify index-exp env)))

(define (simplify:top-level inputs body env)
  `(top-level ,inputs
              ,(simplify body env)))

(define (simplify:outputs names exps env)
  `(outputs ,@(map (lambda (name exp)
                     (list name (simplify exp env)))
                   names exps)))

(define (simplify exp env)
  (match exp
    ((? constant?) exp)
    ((? symbol?)
     (lookup* exp env exp))
    (('if predicate consequent alternate)
     (simplify:if predicate consequent alternate env))
    (('lambda (params ...) body)
     (simplify:lambda params body env))
    (('values exps ...)
     (simplify:values exps env))
    (('let ((names exps) ...) body)
     (simplify:let names exps body env))
    (('let-values ((names exps) ...) body)
     (simplify:let-values names exps body env))
    (('primcall operator args ...)
     (simplify:primcall operator args env))
    (('call operator args ...)
     (simplify:call operator args env))
    (('struct-ref exp field)
     (simplify:struct-ref exp field env))
    (('array-ref array-exp index-exp)
     (simplify:array-ref array-exp index-exp env))
    (('outputs (names exps) ...)
     (simplify:outputs names exps env))
    (('top-level inputs body)
     (simplify:top-level inputs body env))))

(define (simplify* exp)
  (simplify exp (empty-env)))


;;;
;;; Dead code elimination
;;;

;; Find and remove unused variable bindings.  Report errors for unused
;; globals, as they will cause problems later when the graphics
;; driver's GLSL compiler eliminates them.  This also takes care of
;; what would be an issue later on: If an 'outputs' form is bound to
;; an unused variable, the GLSL emitter would emit global variable
;; mutations even though they shouldn't happen!  This is a quirk of
;; 'outputs' being the only form that produces side-effects but dead
;; code elimination takes care of the problem.

(define &seagull-unused-global-error
  (make-exception-type '&seagull-unused-global-error &error '(variable)))

(define make-seagull-unusued-global-error
  (record-constructor &seagull-unused-global-error))

(define seagull-unused-global-variable
  (exception-accessor &seagull-unused-global-error
                      (record-accessor &seagull-unused-global-error 'variable)))

(define (unused-variable? var exp)
  (define (unused-in-list? exps)
    (every (lambda (exp) (unused-variable? var exp)) exps))
  (match exp
    ((? constant?) #t)
    ((? symbol?)
     (not (eq? exp var)))
    (('if predicate consequent alternate)
     (and (unused-variable? var predicate)
          (unused-variable? var consequent)
          (unused-variable? var alternate)))
    (('lambda (params ...) body)
     (unused-variable? var body))
    (('values exps ...)
     (unused-in-list? exps))
    (((or 'let 'let-values) ((names exps) ...) body)
     (and (unused-in-list? exps)
          (unused-variable? var body)))
    (('primcall operator args ...)
     (unused-in-list? args))
    (('call operator args ...)
     (and (unused-variable? var operator)
          (unused-in-list? args)))
    (('struct-ref exp field)
     (unused-variable? var exp))
    (('array-ref array-exp index-exp)
     (and (unused-variable? var array-exp)
          (unused-variable? var index-exp)))
    (('outputs (names exps) ...)
     (and (unused-in-list? names) (unused-in-list? exps)))
    (('top-level _ body)
     (unused-variable? var body))))

(define (prune:list exps)
  (map prune exps))

(define (prune:if predicate consequent alternate)
  `(if ,(prune predicate)
       ,(prune consequent)
       ,(prune alternate)))

(define (prune:lambda params body)
  `(lambda ,params ,(prune body)))

(define (prune:values exps)
  (prune:list exps))

(define (prune:let names exps body)
  (define exps* (prune:list exps))
  (define body* (prune body))
  (define bindings
    (filter-map (lambda (name exp)
                  (if (unused-variable? name body*)
                      #f
                      (list name exp)))
                names exps*))
  ;; Remove 'let' if all bindings are dead.
  (if (null? bindings)
      body*
      `(let ,bindings ,body*)))

(define (prune:let-values names exps body)
  (define bindings
    (filter-map (lambda (names exp)
                  (if (every (lambda (name)
                               (unused-variable? name body))
                             names)
                      #f
                      (list names exp)))
                names exps))
  ;; Remove 'let' if all bindings are dead.
  (if (null? bindings)
      (prune body)
      `(let-values ,bindings ,(prune body))))

(define (prune:primcall operator args)
  `(primcall ,operator ,@(prune:list args)))

(define (prune:call operator args)
  `(call ,(prune operator) ,@(prune:list args)))

(define (prune:struct-ref exp field)
  `(struct-ref ,(prune exp) ,field))

(define (prune:array-ref array-exp index-exp)
  `(array-ref ,(prune array-exp)
              ,(prune index-exp)))

(define (prune:outputs names exps)
  `(outputs ,@(map (lambda (name exp)
                     (list name
                           (prune exp)))
                   names exps)))

(define (prune:top-level qualifiers type-descriptors names body)
  (for-each (lambda (qualifier type-desc name)
              (when (unused-variable? name body)
                (raise-exception
                 (make-exception
                  (make-seagull-unusued-global-error name)
                  (make-exception-with-origin prune:top-level)
                  (make-exception-with-message
                   (format #f "seagull: global variable '~a ~a ~a' is unused"
                           qualifier type-desc name))
                  (make-exception-with-irritants (list name))))))
            qualifiers type-descriptors names)
  `(top-level ,(map list qualifiers type-descriptors names)
              ,(prune body)))

(define (prune exp)
  (match exp
    ((or (? constant?) (? symbol?))
     exp)
    (('if predicate consequent alternate)
     (prune:if predicate consequent alternate))
    (('lambda (params ...) body)
     (prune:lambda params body))
    (('values exps ...)
     (prune:values exps))
    (('let ((names exps) ...) body)
     (prune:let names exps body))
    (('let-values ((names exps) ...) body)
     (prune:let-values names exps body))
    (('primcall operator args ...)
     (prune:primcall operator args))
    (('call operator args ...)
     (prune:call operator args))
    (('struct-ref exp field)
     (prune:struct-ref exp field))
    (('array-ref array-exp index-exp)
     (prune:array-ref array-exp index-exp))
    (('outputs (names exps) ...)
     (prune:outputs names exps))
    (('top-level ((qualifiers type-descs names) ...) body)
     (prune:top-level qualifiers type-descs names body))))


;;;
;;; Function hoisting
;;;

;; Move all lambda bindings to the top-level.  Unfortunately, GLSL
;; does not allow nested function definitions, so nested functions in
;; Seagull only allow free variable references for top-level
;; variables, such as shader inputs and uniforms.

(define &seagull-scope-error
  (make-exception-type '&seagull-scope-error &error '(variable)))

(define make-seagull-scope-error
  (record-constructor &seagull-scope-error))

(define seagull-scope-variable
  (exception-accessor &seagull-scope-error
                      (record-accessor &seagull-scope-error 'variable)))

(define (check-free-variables-in-list exps bound-vars top-level-vars)
  (every (lambda (exp)
           (check-free-variables exp bound-vars top-level-vars))
         exps))

(define (check-free-variables exp bound-vars top-level-vars)
  (match exp
    ((? constant?)
     #t)
    ((? symbol?)
     (or (memq exp bound-vars) ; bound vars: OK
         (memq exp top-level-vars) ; top-level vars: OK
         ;; Free variables that aren't top-level are not allowed because
         ;; GLSL doesn't support closures.
         (raise-exception
          (make-exception
           (make-seagull-scope-error exp)
           (make-exception-with-origin check-free-variables)
           (make-exception-with-message
            "seagull: free variable is not top-level")
           (make-exception-with-irritants (list exp))))))
    (('if predicate consequent alternate)
     (and (check-free-variables predicate bound-vars top-level-vars)
          (check-free-variables consequent bound-vars top-level-vars)
          (check-free-variables alternate bound-vars top-level-vars)))
    (('let ((names exps) ...) body)
     (define bound-vars* (append names bound-vars))
     (and (check-free-variables-in-list exps bound-vars* top-level-vars)
          (check-free-variables body bound-vars* top-level-vars)))
    (('let-values ((names exps) ...) body)
     (define bound-vars* (append (concatenate names) bound-vars))
     (and (check-free-variables-in-list exps bound-vars* top-level-vars)
          (check-free-variables body bound-vars* top-level-vars)))
    (('lambda (params ...) body)
     (check-free-variables body params top-level-vars))
    (('values exps ...)
     (check-free-variables-in-list exps bound-vars top-level-vars))
    ((or ('primcall _ args ...)
         ('call args ...))
     (check-free-variables-in-list args bound-vars top-level-vars))
    (('struct-ref exp _)
     (check-free-variables exp bound-vars top-level-vars))
    (('array-ref array-exp index-exp)
     (and (check-free-variables array-exp bound-vars top-level-vars)
          (check-free-variables index-exp bound-vars top-level-vars)))
    (('outputs (names exps) ...)
     (check-free-variables-in-list exps bound-vars top-level-vars))
    (('top-level ((_ _ names) ...) body)
     (define bound-vars* (append names bound-vars))
     (check-free-variables body bound-vars* top-level-vars))))

(define (hoist:list exps)
  (let-values (((exp-list env-list)
                (unzip2
                 (map (lambda (exp)
                        (call-with-values
                            (lambda ()
                              (hoist-functions exp))
                          list))
                      exps))))
    (values exp-list (apply compose-envs env-list))))

(define (hoist:if predicate consequent alternate)
  (define-values (predicate* predicate-env)
    (hoist-functions predicate))
  (define-values (consequent* consequent-env)
    (hoist-functions consequent))
  (define-values (alternate* alternate-env)
    (hoist-functions alternate))
  (values `(if ,predicate* ,consequent* ,alternate*)
          (compose-envs predicate-env consequent-env alternate-env)))

(define (hoist:let names exps body)
  (define-values (exps* exps-env)
    (hoist:list exps))
  (define-values (body* body-env)
    (hoist-functions body))
  ;; Remove all lambda bindings...
  (define bindings
    (filter-map (lambda (name exp)
                  (match exp
                    (('lambda _ _)
                     #f)
                    (_ (list name exp))))
                names exps*))
  ;; ...and add them to the top-level environment.
  (define env*
    (fold (lambda (name exp env)
            (match exp
              (('lambda _ _)
               (extend-env name exp env))
              (_ env)))
          (compose-envs exps-env body-env)
          names exps*))
  ;; If there are no bindings left, remove the 'let'.
  (values (if (null? bindings)
              body*
              `(let ,bindings ,body*))
          env*))

(define (hoist:let-values names exps body)
  (define-values (exps* exps-env)
    (hoist:list exps))
  (define-values (body* body-env)
    (hoist-functions body))
  ;; Remove all lambda bindings...
  (define bindings
    (filter-map (lambda (names exp)
                  (match names
                    ((name)
                     (match exp
                       (('lambda _ _)
                        #f)
                       (_ (list (list name) exp))))
                    (_
                     (list names exp))))
                names exps*))
  ;; ...and add them to the top-level environment.
  (define env*
    (fold (lambda (name exp env)
            (match names
              ((name)
               (match exp
                 (('lambda _ _)
                  (extend-env name exp env))
                 (_ env)))
              (_ env)))
          (compose-envs exps-env body-env)
          names exps*))
  ;; If there are no bindings left, remove the 'let-values'.
  (values (if (null? bindings)
              body*
              `(let-values ,bindings ,body*))
          env*))

(define (hoist:lambda params body)
  (define-values (body* body-env)
    (hoist-functions body))
  (values `(lambda ,params ,body*) body-env))

(define (hoist:values exps)
  (define-values (exps* exp-env)
    (hoist:list exps))
  (values `(values ,@exps*) exp-env))

(define (hoist:primcall operator args)
  (define-values (args* args-env) (hoist:list args))
  (values `(primcall ,operator ,@args*) args-env))

(define (hoist:call args)
  (define-values (args* args-env) (hoist:list args))
  (values `(call ,@args*) args-env))

(define (hoist:struct-ref exp field)
  (define-values (exp* exp-env) (hoist-functions exp))
  (values `(struct-ref ,exp* ,field) exp-env))

(define (hoist:array-ref array-exp index-exp)
  (define-values (array-exp* array-exp-env)
    (hoist-functions array-exp))
  (define-values (index-exp* index-exp-env)
    (hoist-functions index-exp))
  (values `(array-ref ,array-exp* ,index-exp*)
          (compose-envs array-exp-env index-exp-env)))

(define (hoist:top-level inputs body)
  (define-values (body* body-env)
    (hoist-functions body))
  (values `(top-level ,inputs ,body*)
          body-env))

(define (hoist:outputs names exps)
  (define-values (exps* exp-env)
    (hoist:list exps))
  (values `(outputs
            ,@(map (lambda (name exp)
                     (list name exp))
                   names exps*))
          exp-env))

(define (hoist-functions exp)
  (match exp
    ((or (? constant?) (? symbol?))
     (values exp (empty-env)))
    (('if predicate consequent alternate)
     (hoist:if predicate consequent alternate))
    (('let ((names exps) ...) body)
     (hoist:let names exps body))
    (('let-values ((names exps) ...) body)
     (hoist:let-values names exps body))
    (('lambda (params ...) body)
     (hoist:lambda params body))
    (('values exps ...)
     (hoist:values exps))
    (('primcall operator args ...)
     (hoist:primcall operator args))
    (('call args ...)
     (hoist:call args))
    (('struct-ref exp member)
     (hoist:struct-ref exp member))
    (('array-ref array-exp index-exp)
     (hoist:array-ref array-exp index-exp))
    (('outputs (names exps) ...)
     (hoist:outputs names exps))
    (('top-level inputs body)
     (hoist:top-level inputs body))))

(define (maybe-merge-top-levels new-bindings exp)
  (match exp
    (('top-level bindings body)
     `(top-level ,(append bindings new-bindings) ,body))
    (_
     `(top-level ,new-bindings ,exp))))

(define (hoist-functions* exp)
  (define-values (exp* function-env)
    (hoist-functions exp))
  (define top-level-vars
    (append (env-names function-env)
            (map (match-lambda
                   ((_ _ name) name))
                 (match exp*
                   (('top-level bindings _)
                    bindings)
                   (_ '())))))
  (env-for-each (lambda (name exp)
                  (check-free-variables exp '() top-level-vars))
                function-env)
  (define bindings
    (env-map (lambda (name func)
               `(function ,name ,func))
             function-env))
  (maybe-merge-top-levels bindings exp*))


;;;
;;; Type inference
;;;

;; Walk the expression tree of a type annotated program and solve for
;; all of the type variables using a variant of the Hindley-Milner
;; type inference algorithm extended to handle qualified types (types
;; with predicates.)  GLSL is a statically typed language, but thanks
;; to type inference the user doesn't have to specify any types expect
;; for shader inputs, outputs, and uniforms.

(define &seagull-type-error
  (make-exception-type '&seagull-type-error &error '()))

(define make-seagull-type-error
  (record-constructor &seagull-type-error))

(define (seagull-type-error msg args origin)
  (raise-exception
   (make-exception
    (make-seagull-type-error)
    (make-exception-with-origin origin)
    (make-exception-with-message
     (format #f "seagull type error: ~a" msg))
    (make-exception-with-irritants args))))

(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)
        (struct-type? type)
        (outputs-type? type))
    type)
   ((variable-type? type)
    (if (equal? 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))))
   ((array-type? type)
    (array-type (apply-substitution-to-type (array-type-ref type) from to)
                (array-type-length type)))
   ((type-scheme? type)
    type)
   (else
    (seagull-type-error "invalid type"
                        (list type)
                        apply-substitution-to-type))))

(define (apply-substitutions-to-type type subs)
  (env-fold (lambda (from to type*)
              (apply-substitution-to-type type* from to))
            type
            subs))

(define (apply-substitutions-to-types types subs)
  (map (lambda (type)
         (apply-substitutions-to-type type subs))
       types))

(define (apply-substitution-to-env env from to)
  (env-fold (lambda (name type env*)
              (extend-env name
                          (apply-substitution-to-type type from to)
                          env*))
            (empty-env)
            env))

(define (apply-substitutions-to-env env subs)
  (env-fold (lambda (from to env*)
              (apply-substitution-to-env env* from to))
            env
            subs))

(define (apply-substitutions-to-texp t subs)
  (texp (apply-substitutions-to-types (texp-types t) subs)
        (texp-exp t)))

(define (apply-substitutions-to-exp exp subs)
  (match exp
    ((? type?)
     (apply-substitutions-to-type exp subs))
    ((exps ...)
     (map (lambda (exp)
            (apply-substitutions-to-exp exp subs))
          exps))
    (_ exp)))

;; 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)
    (_ (seagull-type-error "expected single type expression"
                           (list texp)
                           single-type))))

(define (occurs? a b)
  (cond
   ((and (variable-type? a) (variable-type? b))
    (eq? a b))
   ((and (variable-type? 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 (compose-substitutions a b)
  (define b*
    (map (match-lambda
           ((from . to)
            (cons from (apply-substitutions-to-type to a))))
         b))
  (define a*
    (filter-map (match-lambda
                  ((from . to)
                   (if (assq-ref b* from)
                       #f
                       (cons from to))))
                a))
  (append a* b*))

(define* (free-variables-in-type type)
  (cond
   ((or (primitive-type? type)
        (struct-type? type))
    '())
   ((array-type? type)
    (free-variables-in-type (array-type-ref type)))
   ((variable-type? type) (list type))
   ((function-type? 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))))))
   ((type-scheme? type)
    (fold delete
          (free-variables-in-type (type-scheme-ref type))
          (type-scheme-quantifiers type)))
   (else
    (seagull-type-error "unknown type"
                        (list type)
                        free-variables-in-type))))

(define (difference a b)
  (match a
    (() b)
    ((x . rest)
     (if (memq x b)
         (difference rest (delq x b))
         (cons x (difference rest b))))))

(define (free-variables-in-type-scheme type-scheme)
  (difference (type-scheme-quantifiers type-scheme)
              (free-variables-in-type (type-scheme-ref type-scheme))))

(define (free-variables-in-env env)
  (delete-duplicates
   (env-fold (lambda (_name type vars)
               (cond
                ((variable-type? type)
                 (cons (free-variables-in-type type)
                       vars))
                ((type-scheme? type)
                 (cons (free-variables-in-type-scheme type)
                       vars))
                (else vars)))
             '()
             env)))

(define (free-variables-in-predicate pred)
  (match pred
    ((or #t #f) '())
    (((or '= 'substitute) a b)
     (append (free-variables-in-type a)
             (free-variables-in-type b)))
    (((or 'and 'or 'compose) a b)
     (append (free-variables-in-predicate a)
             (free-variables-in-predicate b)))
    (('struct-field struct field var)
     (append (free-variables-in-type struct)
             (free-variables-in-type var)))
    (('array-element array var)
     (append (free-variables-in-type array)
             (free-variables-in-type var)))))

;; Quantified variables are type variables that appear free in the
;; function return types or in the predicate.
(define (generalize type pred env)
  (if (function-type? type)
      (match (difference (delete-duplicates
                          (append (free-variables-in-type type)
                                  (free-variables-in-predicate
                                   (type-predicate-exp pred))))
                         (free-variables-in-env env))
        (() type)
        ((quantifiers ...)
         (type-scheme quantifiers (qualified-type type pred))))
      type))

(define (instantiate type-scheme)
  (define subs
    (fold (lambda (var env)
            (extend-env var (fresh-variable-type) env))
          (empty-env)
          (type-scheme-quantifiers type-scheme)))
  (define type (type-scheme-ref type-scheme))
  (values
   (apply-substitutions-to-type (if (qualified-type? type)
                                    (qualified-type-ref type)
                                    type)
                                subs)
   (if (qualified-type? type)
       (apply-substitutions-to-predicate (qualified-type-predicate type)
                                         subs)
       predicate:succeed)))

(define (maybe-instantiate type)
  (if (type-scheme? type)
      (instantiate type)
      (values type predicate:succeed)))

(define (unify:primitives a b)
  (if (equal? a b)
      '()
      (type-mismatch a b unify:primitives)))

(define (unify:structs a b)
  (if (equal? a b)
      '()
      (type-mismatch a b unify:structs)))

(define (unify:variable a b)
  (cond
   ((eq? a b)
    '())
   ((occurs? a b)
    (seagull-type-error "circular type" (list a b) unify:variable))
   (else
    (list (cons a b)))))

(define (unify:functions a b)
  (define param-subs
    (unify (function-type-parameters a)
           (function-type-parameters b)))
  (define return-subs
    (unify (apply-substitutions-to-types (function-type-returns a)
                                         param-subs)
           (apply-substitutions-to-types (function-type-returns b)
                                         param-subs)))
  (compose-substitutions param-subs return-subs))

(define (unify:lists a rest-a b rest-b)
  (define sub-first (unify a b))
  (define sub-rest
    (unify (apply-substitutions-to-types rest-a sub-first)
           (apply-substitutions-to-types rest-b sub-first)))
  (compose-substitutions sub-first sub-rest))

(define (unify a b)
  (match (list a b)
    (((? primitive-type? a) (? primitive-type? b))
     (unify:primitives a b))
    (((? struct-type? a) (? struct-type? b))
     (unify:structs a b))
    ((or ((? variable-type? a) b)
         (b (? variable-type? a)))
     (unify:variable a b))
    (((? function-type? a) (? function-type? b))
     (unify:functions a b))
    (((? outputs-type?) (? outputs-type?))
     '())
    (((? type?) (? type?))
     (type-mismatch a b unify))
    ((() ())
     '())
    (((a rest-a ...) (b rest-b ...))
     (unify:lists a rest-a b rest-b))
    (_
     (type-mismatch a b unify))))

(define (infer:constant x)
  (values (texp (list (cond
                       ((exact-integer? x)
                        type:int)
                       ((float? x)
                        type:float)
                       ((boolean? x)
                        type:bool)))
                x)
          '()
          predicate:succeed))

(define (infer:variable name env)
  (define-values (type pred)
    (maybe-instantiate (lookup name env)))
  (values (texp (list type) name)
          '()
          pred))

(define (infer:list exps env)
  (let loop ((exps exps)
             (texps '())
             (subs '())
             (pred predicate:succeed))
    (match exps
      (()
       (values (reverse texps) subs pred))
      ((exp . rest)
       (define-values (texp subs* pred*)
         (infer exp env))
       (define-values (new-pred combined-subs)
         (eval-predicate* (predicate:compose pred pred*)
                          (compose-substitutions subs subs*)))
       (loop rest
             (cons texp texps)
             combined-subs
             new-pred)))))

(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 predicate env))
  (define predicate-unify-subs
    (unify (texp-types predicate-texp) (list type:bool)))
  ;; Combine the substitutions and apply them to the environment.
  (define combined-subs-0
    (compose-substitutions predicate-subs predicate-unify-subs))
  (define env0
    (apply-substitutions-to-env env combined-subs-0))
  ;; 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 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 alternate env1))
  (define combined-subs-2
    (compose-substitutions combined-subs-1 alternate-subs))
  ;; Eval combined predicate.
  (define-values (pred combined-subs-3)
    (eval-predicate* (predicate:compose predicate-pred
                                        consequent-pred
                                        alternate-pred)
                     combined-subs-2))
  ;; ;; Apply final set of substitutions to the types of both branches.
  (define consequent-texp*
    (apply-substitutions-to-texp consequent-texp combined-subs-3))
  (define alternate-texp*
    (apply-substitutions-to-texp alternate-texp combined-subs-3))
  (values (texp (texp-types consequent-texp)
                `(if ,predicate-texp ,consequent-texp ,alternate-texp))
          combined-subs-3
          pred))

(define (infer:lambda params body env)
  ;; Each function parameter gets a fresh type variable.
  (define param-types (fresh-variable-types-for-list params))
  ;; The type environment is extended with the function parameters.
  (define env*
    (fold (lambda (param type env*)
            (extend-env param type env*))
          env params param-types))
  (define-values (body* body-subs body-pred)
    (infer body env*))
  (define-values (pred subs)
    (eval-predicate* body-pred body-subs))
  (values (texp (list (generalize
                       (function-type (apply-substitutions-to-types param-types
                                                                    subs)
                                      (texp-types body*))
                       pred env))
                `(lambda ,params ,body*))
          subs predicate:succeed))

(define (check-arity type arity)
  (define (arity-error)
    (seagull-type-error "wrong number of arguments"
                        (list type arity) check-arity))
  (cond
   ((function-type? type)
    (if (= (length (function-type-parameters type)) arity)
        type
        (arity-error)))
   ((function-case-type? type)
    (let ((function (function-case-type-ref type arity)))
      (or function (arity-error))))
   ;; TODO: We aren't actually checking arity here.
   ((type-scheme? type)
    type)
   ((type? type)
    (seagull-type-error "expected a function" (list type) check-arity))))

(define (infer:primitive-call operator args env)
  (define primitive (lookup-seagull-primitive operator))
  ;; Primitive functions may be overloaded and need to be instantiated
  ;; with fresh type variables.
  (define-values (operator-type operator-pred)
    (maybe-instantiate
     (check-arity (seagull-primitive-type primitive)
                  (length args))))
  ;; Infer the arguments.
  (define-values (args* arg-subs arg-pred)
    (infer:list args env))
  ;; Generate fresh type variables to unify against the return types
  ;; of the operator.
  (define return-vars
    (fresh-variable-types-for-list (function-type-returns operator-type)))
  (define call-subs
    (unify operator-type
           (function-type (map single-type args*)
                          return-vars)))
  ;; Apply substitutions to the predicate and then eval it, producing
  ;; a simplified predicate and a set of substitutions.
  (define-values (pred combined-subs)
    (eval-predicate* (predicate:compose operator-pred arg-pred)
                     (compose-substitutions arg-subs call-subs)))
  (values (texp (apply-substitutions-to-types return-vars combined-subs)
                `(primcall ,operator
                           ,@(map (lambda (arg)
                                    (apply-substitutions-to-texp arg
                                                                 combined-subs))
                                  args*)))
          combined-subs
          pred))

(define (infer:call operator args env)
  ;; The type signature of primitive functions can be looked up
  ;; directly in the environment.
  (define-values (operator* operator-subs operator-pred)
    (infer operator env))
  (define env*
    (apply-substitutions-to-env env operator-subs))
  ;; Infer the arguments.
  (define-values (args* arg-subs arg-pred)
    (infer:list args env*))
  (define combined-subs-0
    (compose-substitutions operator-subs arg-subs))
  ;; Generate fresh type variables to unify against the return types
  ;; of the operator.
  (define operator-type (single-type operator*))
  (define return-vars
    (fresh-variable-types-for-list
     (function-type-returns operator-type)))
  (define call-subs
    (unify (apply-substitutions-to-type operator-type combined-subs-0)
           (function-type (apply-substitutions-to-types (map single-type args*)
                                                        combined-subs-0)
                          return-vars)))
  ;; Eval predicate.
  (define-values (pred combined-subs)
    (eval-predicate* (predicate:compose operator-pred
                                        arg-pred)
                     (compose-substitutions combined-subs-0 call-subs)))
  (values (texp (apply-substitutions-to-types return-vars combined-subs)
                `(call ,(apply-substitutions-to-texp operator* combined-subs)
                       ,@(map (lambda (arg)
                                (apply-substitutions-to-texp arg
                                                             combined-subs))
                              args*)))
          combined-subs
          pred))

(define (infer:struct-ref exp field env)
  (define-values (exp* exp-subs exp-pred)
    (infer exp env))
  (define exp-type (single-type exp*))
  (define tvar (fresh-variable-type))
  (define-values (pred combined-subs)
    (eval-predicate* (predicate:compose (predicate:struct-field exp-type field tvar)
                                        exp-pred)
                     exp-subs))
  (values (texp (list (apply-substitutions-to-type tvar combined-subs))
                `(struct-ref ,(apply-substitutions-to-texp exp* combined-subs)
                             ,field))
          combined-subs
          pred))

(define (infer:array-ref array-exp index-exp env)
  (define-values (array-exp* array-exp-subs array-exp-pred)
    (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 index-exp env*))
  (define index-type (single-type index-exp*))
  (define combined-subs
    (compose-substitutions array-exp-subs index-exp-subs))
  ;; Array indices must be integers.
  (define unify-subs
    (unify (apply-substitutions-to-type index-type combined-subs) type:int))
  (define tvar (fresh-variable-type))
  (define-values (pred subs)
    (eval-predicate* (predicate:compose (predicate:array-element array-type tvar)
                                         array-exp-pred
                                         index-exp-pred)
                     (compose-substitutions combined-subs unify-subs)))
  (define array-exp**
    (apply-substitutions-to-texp array-exp* subs))
  (define index-exp**
    (apply-substitutions-to-texp index-exp* subs))
  (values (texp (list tvar)
                `(array-ref ,array-exp** ,index-exp**))
          subs
          pred))

(define (infer:let names exps body env)
  (define-values (exps* exp-subs exp-pred)
    (infer:list exps env))
  (define exp-types (map single-type exps*))
  (define env*
    (fold extend-env
          (apply-substitutions-to-env env exp-subs)
          names
          exp-types))
  (define-values (body* body-subs body-pred)
    (infer body env*))
  (define-values (pred combined-subs)
    (eval-predicate* (predicate:compose exp-pred body-pred)
                     (compose-substitutions exp-subs body-subs)))
  (define bindings
    (map (lambda (name exp)
           (let ((num-types (length (texp-types exp))))
             (unless (= num-types 1)
               (seagull-type-error (format #f "expected 1 value, got ~a"
                                           num-types)
                                   (list name exp)
                                   infer:let))
            (list name (apply-substitutions-to-texp exp combined-subs))))
         names exps*))
  (values (texp (texp-types body*)
                `(let ,bindings
                   ,(apply-substitutions-to-texp body* combined-subs)))
          combined-subs
          pred))

(define (infer:let-values names exps body env)
  (define-values (exps* exp-subs exp-pred)
    (infer:list exps env))
  (define exp-types (map texp-types exps*))
  (define env*
    (fold (lambda (names types env)
            (fold extend-env env names types))
          (apply-substitutions-to-env env exp-subs)
          names
          exp-types))
  (define-values (body* body-subs body-pred)
    (infer body env*))
  (define-values (pred combined-subs)
    (eval-predicate* (predicate:compose exp-pred body-pred)
                     (compose-substitutions exp-subs body-subs)))
  (define bindings
    (map (lambda (names exp)
           (let ((num-names (length names))
                 (num-types (length (texp-types exp))))
             (unless (= num-names num-types)
               (seagull-type-error (format #f "expected ~a ~a, got ~a"
                                           num-names
                                           (if (= num-names 1) "value" "values")
                                           num-types)
                                   (list names exp)
                                   infer:let-values))
             (list names
                   (apply-substitutions-to-texp exp combined-subs))))
         names exps*))
  (values (texp (texp-types body*)
                `(let-values ,bindings
                   ,(apply-substitutions-to-texp body* combined-subs)))
          combined-subs
          pred))

(define (infer:values exps env)
  (define-values (exps* exp-subs exp-pred)
    (infer:list exps env))
  (values (texp (map single-type exps*)
                `(values ,@exps*))
          exp-subs
          exp-pred))

(define (infer:outputs names exps env)
  (define-values (exps* exp-subs exp-pred)
    (infer:list exps env))
  (define unify-subs
    (unify (map single-type exps*)
           (map (lambda (name)
                  (lookup name env))
                names)))
  ;; Eval predicate.
  (define-values (pred combined-subs)
    (eval-predicate* exp-pred (compose-substitutions exp-subs unify-subs)))
  (values (texp (list type:outputs)
                `(outputs
                  ,@(map (lambda (name exp)
                           (list name (apply-substitutions-to-texp
                                       exp combined-subs)))
                         names exps*)))
          combined-subs
          pred))

(define (infer:top-level bindings body env)
  (define (infer-bindings bindings texps subs pred env)
    (match bindings
      (()
       (values (reverse texps) subs pred env))
      ((('function name exp) . rest)
       (define-values (texp subs* pred*)
         (infer exp env))
       (define-values (new-pred combined-subs)
         (eval-predicate* (predicate:compose pred pred*)
                          (compose-substitutions subs subs*)))
       (define env*
         (apply-substitutions-to-env (extend-env name (single-type texp) env)
                                     combined-subs))
       (infer-bindings rest
                       (cons texp texps)
                       combined-subs
                       new-pred
                       env*))
      (((_ desc name) . rest)
       (define type (type-descriptor->type desc))
       (infer-bindings rest
                       (cons (list type) texps)
                       subs
                       pred
                       (extend-env name type env)))))
  (define qualifiers (map first bindings))
  (define names
    (map (match-lambda
           (('function name _) name)
           ((_ _ name) name))
         bindings))
  (define type-names
    (map (match-lambda
           (((? top-level-qualifier?) type-name _) type-name)
           (_ #f))
         bindings))
  (define-values (exps exp-subs exp-pred env*)
    (infer-bindings bindings '() '() predicate:succeed env))
  (define-values (body* body-subs body-pred)
    (infer body env*))
  (define-values (pred combined-subs)
    (eval-predicate* (predicate:compose exp-pred body-pred)
                     (compose-substitutions exp-subs body-subs)))
  (define bindings*
    (map (match-lambda*
           (((? top-level-qualifier? qualifier) type-name name _)
            (list qualifier type-name name))
           (('function _ name exp)
            `(function ,name ,(apply-substitutions-to-exp exp combined-subs))))
         qualifiers type-names names exps))
  (values (texp (texp-types body*)
                `(top-level ,bindings* ,body*))
          combined-subs
          pred))

;; Inference returns 3 values:
;; - a typed expression
;; - a list of substitutions
;; - a type predicate
(define (infer exp env)
  (match exp
    ((? constant?)
     (infer:constant exp))
    ((? symbol? name)
     (infer:variable name env))
    (('if predicate consequent alternate)
     (infer:if predicate consequent alternate env))
    (('let ((names exps) ...) body)
     (infer:let names exps body env))
    (('let-values ((names exps) ...) body)
     (infer:let-values names exps body env))
    (('lambda (params ...) body)
     (infer:lambda params body env))
    (('values exps ...)
     (infer:values exps env))
    (('primcall operator args ...)
     (infer:primitive-call operator args env))
    (('call operator args ...)
     (infer:call operator args env))
    (('struct-ref exp field)
     (infer:struct-ref exp field env))
    (('array-ref array-exp index-exp)
     (infer:array-ref array-exp index-exp env))
    (('outputs (names exps) ...)
     (infer:outputs names exps env))
    (('top-level bindings body)
     (infer:top-level bindings body env))
    ;; User code shouldn't trigger this, only us screwing up an
    ;; earlier compiler pass.
    (_ (error "unknown form" exp))))

(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* exp stage)
  (infer exp (infer:top-level-env stage)))


;;;
;;; Overloaded functions
;;;

;; Replace quantified functions ('type-scheme' expressions) with a series
;; of non-quantified function type specifications, one for each unique
;; type of call in the program.

(define (find-signatures:list name texps)
  (append-map (lambda (texp)
                (find-signatures name texp))
              texps))

(define (find-signatures:if name predicate consequent alternate)
  (append (find-signatures name predicate)
          (find-signatures name consequent)
          (find-signatures name alternate)))

(define (find-signatures:let name binding-texps body)
  (append (find-signatures:list name binding-texps)
          (find-signatures name body)))

(define (find-signatures:array-ref name array index)
  (append (find-signatures name array)
          (find-signatures name index)))

(define (find-signatures name texp)
  (match (texp-exp texp)
    ((or (? constant?) (? symbol?))
     '())
    (('if predicate consequent alternate)
     (find-signatures:if name predicate consequent alternate))
    (('let ((_ exps) ...) body)
     (find-signatures:let name exps body))
    (('values exps ...)
     (find-signatures:list name exps))
    (('primcall _ args ...)
     (find-signatures:list name args))
    (('call operator args ...)
     (cons (if (eq? (texp-exp operator) name)
               (function-type (map single-type args)
                              (texp-types texp)))
           (find-signatures:list name args)))
    (('struct-ref struct _)
     (find-signatures name struct))
    (('array-ref array index)
     (find-signatures:array-ref name array index))
    (('outputs (_ exps) ...)
     (find-signatures:list name exps))))

(define (vars->subs exp env)
  (match exp
    (('t ((? variable-type? tvar)) (? symbol? name))
     (let ((type (lookup* name env)))
       (if type
           (list (cons tvar type))
           '())))
    ((head . rest)
     (delete-duplicates
      (append (vars->subs head env)
              (vars->subs rest env))))
    (_ '())))

(define (untype x)
  (match x
    (('t (_ ...) exp)
     (untype exp))
    ((exp . rest)
     (cons (untype exp) (untype rest)))
    (_ x)))

(define (resolve-overloads program stage)
  ;; Find all of the struct types used in the program.  They will be
  ;; used to generate overloaded functions that take one or more
  ;; structs as arguments.
  ;;(define structs (delete-duplicates (find-structs program)))
  (match program
    (('t types ('top-level bindings body))
     (define bindings*
       (let loop ((bindings bindings)
                  (globals (empty-env)))
         (match bindings
           (() '())
           ((('function name ('t ((? type-scheme? type)) func)) . rest)
            (define qtype (type-scheme-ref type))
            (define func-type (qualified-type-ref qtype))
            (append (map (lambda (call-type)
                           (define subs
                             (unify func-type call-type))
                           (define type*
                             (apply-substitutions-to-type func-type subs))
                           (define params
                             (match func
                               (('lambda (params ...) _)
                                params)))
                           (define env
                             (compose-envs (fold extend-env
                                                 (empty-env)
                                                 params
                                                 (function-type-parameters
                                                  type*))
                                           globals))
                           (match func
                             (('lambda _ body)
                              (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*
                             (apply-substitutions-to-exp func subs*))
                           `(function ,name (t (,type*) ,func*)))
                         (delete-duplicates
                          (find-signatures name body)))
                    (loop rest globals)))
           ((('function name texp) . rest)
            (cons `(function ,name ,texp)
                  (loop rest globals)))
           (((qualifier type name) . rest)
            (cons (list qualifier type name)
                  (loop rest
                        (extend-env name
                                    (type-descriptor->type type)
                                    globals)))))))
     `(t ,types (top-level ,bindings* ,body)))))


;;;
;;; GLSL emission
;;;

;; Transform a fully typed Seagull program into a string of GLSL code.

(define (type-descriptor->glsl desc)
  (match desc
    ((? symbol?)
     (match (lookup-type desc)
       ((? primitive-type? primitive)
        (primitive-type-glsl-name primitive))
       ((? struct-type? struct)
        (struct-type-glsl-name struct))))
    (('array desc* length)
     (format #f "~a[~a]"
             (type-descriptor->glsl desc*)
             length))))

(define (type->glsl type)
  (cond
   ((primitive-type? type)
    (primitive-type-glsl-name type))
   ((struct-type? type)
    (struct-type-glsl-name type))
   ((array-type? type)
    (format #f "~a[~a]"
            (type->glsl (array-type-ref type))
            length))))

(define (single-temp temps)
  (match temps
    ((temp) temp)))

(define (indent n port)
  (when (> n 0)
    (display (make-string (* n 2) #\space) port)))

(define (emit:int n stage version port level)
  (define temp (unique-identifier))
  (indent level port)
  (format port "int ~a = ~a;\n" temp n)
  (list temp))

(define (emit:float n stage version port level)
  (define temp (unique-identifier))
  (indent level port)
  (format port "float ~a = ~a;\n" temp
          (if (inf? n) "1.0 / 0.0" n))
  (list temp))

(define (emit:boolean b stage version port level)
  (define temp (unique-identifier))
  (indent level port)
  (format port "bool ~a = ~a;\n" temp (if b "true" "false"))
  (list temp))

(define (emit:declaration type lhs rhs port level)
  (unless (outputs-type? type)
    (indent level port)
    (if rhs
        (format port "~a ~a = ~a;\n" (type->glsl type) lhs rhs)
        (format port "~a ~a;\n" (type->glsl type) lhs))))

(define (emit:declarations types lhs-list rhs-list port level)
  (define rhs-list* (if rhs-list rhs-list (make-list (length lhs-list) #f)))
  (for-each (lambda (type lhs rhs)
              (emit:declaration type lhs rhs port level))
            types lhs-list rhs-list*))

(define (emit:mov a b port level)
  (when a
    (indent level port)
    (format port "~a = ~a;\n" a b)))

(define (emit:function name type params body stage version port level)
  (define param-types (function-type-parameters type))
  (define return-types (function-type-returns type))
  (define outputs (unique-identifiers-for-list return-types))
  (indent level port)
  (format port "void ~a(" name)
  (let loop ((params (append (zip (make-list (length params) 'in)
                                  param-types
                                  params)
                             (zip (make-list (length return-types) 'out)
                                  return-types
                                  outputs)))
             (first? #t))
    (match params
      (() #t)
      (((qualifier type name) . rest)
       (unless first?
         (display ", " port))
       (format port "~a ~a ~a"
               qualifier (type->glsl type) name)
       (loop rest #f))))
  (display ") {\n" port)
  (define body-temps (emit-glsl body stage version port (+ level 1)))
  (for-each (lambda (output temp)
              (emit:mov output temp port (+ level 1)))
            outputs body-temps)
  (indent level port)
  (display "}\n" port))

(define (emit:if predicate consequent alternate stage version port level)
  (define if-temps
    (if (equal? (texp-types consequent) (list type:outputs))
        '(#f)
        (unique-identifiers-for-list (texp-types consequent))))
  (emit:declarations (texp-types consequent) if-temps #f port level)
  (define predicate-temp
    (single-temp (emit-glsl predicate stage version port level)))
  (indent level port)
  (format port "if(~a) {\n" predicate-temp)
  (define consequent-temps
    (emit-glsl consequent stage version port (+ level 1)))
  (for-each (lambda (lhs rhs)
              (emit:mov lhs rhs port (+ level 1)))
            if-temps consequent-temps)
  (indent level port)
  (display "} else {\n" port)
  (define alternate-temps
    (emit-glsl alternate stage version port (+ level 1)))
  (for-each (lambda (lhs rhs)
              (emit:mov lhs rhs port (+ level 1)))
            if-temps alternate-temps)
  (indent level port)
  (display "}\n" port)
  if-temps)

(define (emit:values exps stage version port level)
  (append-map (lambda (exp)
                (emit-glsl exp stage version port level))
              exps))

(define (emit:let types names exps body stage version port level)
  (define binding-temps
    (map (lambda (exp)
           (single-temp (emit-glsl exp stage version port level)))
         exps))
  (define binding-types (map single-type exps))
  (emit:declarations binding-types names binding-temps port level)
  (define body-temps (emit-glsl body stage version port level))
  (define let-temps (unique-identifiers-for-list types))
  (emit:declarations (texp-types body) let-temps body-temps port level)
  let-temps)

(define (emit:let-values types names exps body stage version port level)
  (define names* (concatenate names))
  (define binding-temps
    (append-map (lambda (exp)
                  (emit-glsl exp stage version port level))
                exps))
  (define binding-types (append-map texp-types exps))
  (emit:declarations binding-types names* binding-temps port level)
  (define body-temps (emit-glsl body stage version port level))
  (define let-temps (unique-identifiers-for-list types))
  (emit:declarations (texp-types body) let-temps body-temps port level)
  let-temps)

(define (emit:primcall type operator args stage version port level)
  (define primitive (lookup-seagull-primitive operator))
  (define operator* (seagull-primitive-glsl-name primitive))
  (define arg-temps
    (map (lambda (arg)
           (single-temp (emit-glsl arg stage version port level)))
         args))
  (define output-temp (unique-identifier))
  (indent level port)
  (format port "~a ~a = "
          (type->glsl type)
          output-temp)
  ((seagull-primitive-emit primitive) arg-temps port)
  (format port ";\n")
  (list output-temp))

(define (emit:call types operator args stage version port level)
  (define operator-name (single-temp (emit-glsl operator stage version port)))
  (define arg-temps
    (map (lambda (arg)
           (single-temp (emit-glsl arg stage version port level)))
         args))
  (define output-temps (unique-identifiers-for-list types))
  (emit:declarations types output-temps #f port level)
  (indent level port)
  (format port "~a(~a);\n"
          operator-name
          (string-join (map symbol->string (append arg-temps output-temps))
                       ", "))
  output-temps)

(define (emit:struct-ref type exp field stage version port level)
  (define input-temp (single-temp (emit-glsl exp stage version port level)))
  (define output-temp (unique-identifier))
  (indent level port)
  (format port "~a ~a = ~a.~a;\n"
          (type->glsl type)
          output-temp
          input-temp
          field)
  (list output-temp))

(define (emit:array-ref type array-exp index-exp stage version port level)
  (define array-temp (single-temp (emit-glsl array-exp stage version port level)))
  (define index-temp (single-temp (emit-glsl index-exp stage version port level)))
  (define output-temp (unique-identifier))
  (indent level port)
  (format port "~a ~a = ~a[~a];\n"
          (type->glsl type)
          output-temp
          array-temp
          index-temp)
  (list output-temp))

(define (emit:top-level bindings body stage version port level)
  (define (glsl-qualifier qualifier)
    (case qualifier
      ((in)
       (if (string>= version "1.3") 'in 'attribute))
      ((out)
       (if (string>= version "1.3") 'out 'varying))
      ((uniform)
       'uniform)))
  (for-each (match-lambda
              (((? top-level-qualifier? qualifier) type-desc name)
               (format port "~a ~a ~a;\n"
                       (glsl-qualifier qualifier)
                       (type-descriptor->glsl type-desc)
                       name))
              (('function name ('t (type) ('lambda params body)))
               (emit:function name type params body stage version port level)))
            bindings)
  (display "void main() {\n" port)
  (emit-glsl body stage version port (+ level 1))
  (display "}\n" port))

(define (glsl-output-name name)
  (let ((variable (lookup-output-variable name)))
    (if (seagull-variable? variable)
        (seagull-variable-glsl-name variable)
        name)))

(define (emit:outputs names exps stage version port level)
  (if (and (eq? stage 'fragment) (null? names))
      (begin
        (indent level port)
        (format port "discard;\n"))
      (for-each (lambda (name exp)
                  (match (emit-glsl exp stage version port level)
                    ((temp)
                     (indent level port)
                     (format port "~a = ~a;\n"
                             (glsl-output-name name)
                             temp))))
                names exps))
  '(#f))

(define (glsl-input-name name)
  (let ((variable (lookup-input-variable name)))
    (if (seagull-variable? variable)
        (seagull-variable-glsl-name variable)
        name)))

(define* (emit-glsl exp stage version port #:optional (level 0))
  (match exp
    (('t _ (? exact-integer? n))
     (emit:int n stage version port level))
    (('t _ (? float? n))
     (emit:float n stage version port level))
    (('t _ (? boolean? b))
     (emit:boolean b stage version port level))
    (('t _ (? symbol? var))
     (list (glsl-input-name var)))
    (('t _ ('if predicate consequent alternate))
     (emit:if predicate consequent alternate stage version port level))
    (('t _ ('values exps ...))
     (emit:values exps stage version port level))
    (('t types ('let ((names exps) ...) body))
     (emit:let types names exps body stage version port level))
    (('t types ('let-values ((names exps) ...) body))
     (emit:let-values types names exps body stage version port level))
    (('t (type) ('primcall op args ...))
     (emit:primcall type op args stage version port level))
    (('t types ('call operator args ...))
     (emit:call types operator args stage version port level))
    (('t (type) ('struct-ref exp field))
     (emit:struct-ref type exp field stage version port level))
    (('t (type) ('array-ref array-exp index-exp))
     (emit:array-ref type array-exp index-exp stage version port level))
    (('t _ ('outputs (names exps) ...))
     (emit:outputs names exps stage version port level))
    (('t _ ('top-level (bindings ...) body))
     (emit:top-level bindings body stage version port level))))


;;;
;;; Compiler front-end
;;;

;; Combine all of the compiler passes on a user provided program and
;; emit GLSL code if the program is valid.

(define &seagull-compiler-error
  (make-exception-type '&seagull-compiler-error &error '()))

(define make-seagull-compiler-error
  (record-constructor &seagull-compiler-error))

(define (seagull-compiler-error msg args origin)
  (raise-exception
   (make-exception
    (make-seagull-compiler-error)
    (make-exception-with-origin origin)
    (make-exception-with-message
     (format #f "seagull compilation error: ~a" msg))
    (make-exception-with-irritants args))))

(define-record-type <seagull-global>
  (make-seagull-global qualifier type-descriptor name)
  seagull-global?
  (qualifier seagull-global-qualifier)
  (type-descriptor seagull-global-type-descriptor)
  (name seagull-global-name))

(define-record-type <seagull-module>
  (%make-seagull-module stage inputs outputs uniforms source compiled
                        global-map max-id)
  seagull-module?
  (stage seagull-module-stage)
  (inputs seagull-module-inputs)
  (outputs seagull-module-outputs)
  (uniforms seagull-module-uniforms)
  (source seagull-module-source)
  (compiled seagull-module-compiled)
  ;; Original name -> alpha converted name mapping for inputs,
  ;; outputs, and uniforms.
  (global-map seagull-module-global-map)
  (max-id seagull-module-max-id))

(define* (make-seagull-module #:key stage inputs outputs uniforms source
                              compiled global-map max-id)
  (%make-seagull-module stage inputs outputs uniforms source compiled
                        global-map max-id))

(define (seagull-module-vertex? module)
  (eq? (seagull-module-stage module) 'vertex))

(define (seagull-module-fragment? module)
  (eq? (seagull-module-stage module) 'fragment))

(define* (compile-seagull #:key stage body
                          (inputs '()) (outputs '()) (uniforms '()))
  (unless (memq stage '(vertex fragment))
    (seagull-compiler-error "invalid shader stage" (list stage) compile-seagull))
  (parameterize ((unique-identifier-counter 0)
                 (unique-variable-type-counter 0))
    (let ((source* `(top-level ,(append inputs outputs uniforms)
                               ,body)))
      (define-values (expanded global-map)
        (expand* source* stage))
      (let* ((simplified (simplify expanded (empty-env)))
             (pruned (prune simplified))
             (hoisted (hoist-functions* pruned))
             (inferred (infer* hoisted stage))
             (resolved (resolve-overloads inferred stage)))
        (values resolved global-map (unique-identifier-counter))))))

(define (specs->globals specs)
  (map (match-lambda
         ((qualifier type-desc name)
          (make-seagull-global qualifier type-desc name)))
       specs))

(define (partition-globals exps)
  (let loop ((exps exps)
             (inputs '())
             (outputs '())
             (uniforms '()))
    (match exps
      (((and ('in _ _) spec) . rest)
       (loop rest (cons spec inputs) outputs uniforms))
      (((and ('out _ _) spec) . rest)
       (loop rest inputs (cons spec outputs) uniforms))
      (((and ('uniform _ _) spec) . rest)
       (loop rest inputs outputs (cons spec uniforms)))
      ((body ...)
       (values (reverse inputs)
               (reverse outputs)
               (reverse uniforms)
               `(begin ,@body))))))

;; Allow importing Scheme values into Seagull expressions with special
;; '$' syntax, such as numeric constants or user-defined shader types.
(define-syntax seagull-quasiquote
  (syntax-rules ($)
    ((_ ($ x)) x)
    ((_ (x ...)) (list (seagull-quasiquote x) ...))
    ((_ x) (quote x))))

;; Using syntax-case allows us to compile shaders to their fully typed
;; intermediate form at compile time, leaving only GLSL emission for
;; runtime.
(define-syntax define-shader-stage
  (lambda (x)
    (syntax-case x ()
      ((_ name stage body ...)
       (let*-values (((inputs outputs uniforms body)
                      (partition-globals
                       (syntax->datum
                        (eval #'(seagull-quasiquote (body ...))
                              (current-module))))))
         (define-values (compiled global-map max-id)
           (compile-seagull #:stage (syntax->datum #'stage)
                            #:inputs inputs
                            #:outputs outputs
                            #:uniforms uniforms
                            #:body body))
         (with-syntax ((inputs (datum->syntax x inputs))
                       (outputs (datum->syntax x outputs))
                       (uniforms (datum->syntax x uniforms))
                       (compiled (datum->syntax x compiled))
                       (global-map (datum->syntax x global-map))
                       (max-id (datum->syntax x max-id))
                       (body (datum->syntax x body)))
           #'(define name
               (make-seagull-module #:stage 'stage
                                    #:inputs (specs->globals 'inputs)
                                    #:outputs (specs->globals 'outputs)
                                    #:uniforms (specs->globals 'uniforms)
                                    #:source 'body
                                    #:compiled 'compiled
                                    #:global-map 'global-map
                                    #:max-id max-id))))))))

(define-syntax-rule (define-vertex-shader name specs source ...)
  (define-shader-stage name vertex specs source ...))

(define-syntax-rule (define-fragment-shader name specs source ...)
  (define-shader-stage name fragment specs source ...))

(define (vertex-outputs-match-fragment-inputs? vertex fragment)
  (let ((fragment-inputs (seagull-module-inputs fragment)))
    (every (lambda (o1)
             (any (lambda (o2)
                    (and (eq? (seagull-global-name o1)
                              (seagull-global-name o2))
                         (equal? (seagull-global-type-descriptor o1)
                                 (seagull-global-type-descriptor o2))))
                  fragment-inputs))
           (seagull-module-outputs vertex))))

(define (uniforms-compatible? vertex fragment)
  (let ((fragment-uniforms (seagull-module-uniforms fragment)))
    (every (lambda (u1)
             (every (lambda (u2)
                      (if (eq? (seagull-global-name u1)
                               (seagull-global-name u2))
                          (equal? (seagull-global-type-descriptor u1)
                                  (seagull-global-type-descriptor u2))
                          #t))
                    fragment-uniforms))
           (seagull-module-outputs vertex))))

(define (rewrite-variables exp subs)
  (match exp
    ((? symbol?)
     (or (assq-ref subs exp) exp))
    (() '())
    ((exp* . rest)
     (cons (rewrite-variables exp* subs)
           (rewrite-variables rest subs)))
    (_ exp)))

(define (link-vertex-outputs-with-fragment-inputs vertex fragment)
  (define (map-globals specs global-map)
    (map (lambda (global)
           (let ((name (seagull-global-name global)))
             (cons name (assq-ref global-map name))))
         specs))
  (define (alpha-rename name-map)
    (map (match-lambda
           ((original-name . alpha-name)
            (cons alpha-name (unique-identifier))))
         name-map))
  (define (remap specs global-map alpha-map)
    (map (lambda (global)
           (let ((name (seagull-global-name global)))
            (cons (assq-ref alpha-map (assq-ref global-map name))
                  name)))
         specs))
  (let* ((vertex-global-map (seagull-module-global-map vertex))
         ;; Create a Scheme name -> alpha-converted GLSL name mapping
         ;; for vertex outputs.
         (vertex-output-map (map-globals (seagull-module-outputs vertex)
                                         vertex-global-map))
         ;; Create a Scheme name -> alpha-converted GLSL name mapping
         ;; for vertex uniforms.
         (vertex-uniform-map (map-globals (seagull-module-uniforms vertex)
                                          vertex-global-map))
         ;; Give new GLSL names to the vertex outputs and uniforms
         ;; that are unique to both the vertex and fragment shaders.
         ;; The vertex output names are changed so that the fragment
         ;; input names can be changed to match.  The vertex uniform
         ;; names are changed so that the names do not clash with
         ;; fragment globals.
         (vertex-output-alpha-map (alpha-rename vertex-output-map))
         (vertex-uniform-alpha-map (alpha-rename vertex-uniform-map))
         (fragment-global-map (seagull-module-global-map fragment))
         ;; Create a Scheme name -> alpha-converted GLSL name mapping
         ;; for fragment inputs.
         (fragment-input-map (map-globals (seagull-module-inputs fragment)
                                          fragment-global-map))
         ;; Create a Scheme name -> alpha-converted GLSL name mapping
         ;; for fragment uniforms.
         (fragment-uniform-map (map-globals (seagull-module-uniforms fragment)
                                            fragment-global-map))
         ;; Give new names to the fragment uniforms so that the names
         ;; do not clash with vertex globals and also that any
         ;; uniforms in the vertex shader have the *same* name in the
         ;; fragment shader.
         (fragment-uniform-alpha-map
          (map (match-lambda
                 ((original-name . alpha-name)
                  (cons alpha-name
                        (or (assq-ref vertex-uniform-alpha-map
                                      (assq-ref vertex-uniform-map original-name))
                            (unique-identifier)))))
               fragment-uniform-map))
         ;; This one is a little messy but what's happening is that
         ;; the GLSL name for each fragment output is mapped to the
         ;; respective renamed input.  Vertex shader output names must
         ;; match fragment shader input names.
         (fragment-input-alpha-map
          (append (map (lambda (input)
                         (let ((name (seagull-global-name input)))
                           (cons (assq-ref fragment-global-map
                                           name)
                                 (assq-ref vertex-output-alpha-map
                                           (assq-ref vertex-global-map
                                                     name)))))
                       (seagull-module-inputs fragment)))))
    ;; Rewrite the intermediate compiled forms of both shader stages
    ;; to replace global variable names as needed.
    (values (rewrite-variables (seagull-module-compiled vertex)
                               (append vertex-uniform-alpha-map
                                       vertex-output-alpha-map))
            (rewrite-variables (seagull-module-compiled fragment)
                               (append fragment-uniform-alpha-map
                                       fragment-input-alpha-map))
            ;; Generate a list of alpha-converted GLSL name -> Scheme
            ;; name mappings.  This will be given to the OpenGL shader
            ;; constructor to map the human readable uniform names to
            ;; the names they've been given by the compiler.
            (append (remap (seagull-module-uniforms vertex)
                           vertex-global-map
                           vertex-uniform-alpha-map)
                    (remap (seagull-module-uniforms fragment)
                           fragment-global-map
                           fragment-uniform-alpha-map)))))

(define (seagull-module-uniform-map module)
  (let ((global-map (seagull-module-global-map module)))
    (map (match-lambda
           ((_ _ name)
            (cons (assq-ref global-map name) name)))
         (seagull-module-uniforms module))))

(define (emit-version-preprocessor version port)
  (cond
   ((string>= version "3.3")
    (format port "#version 330\n"))
   ((string>= version "1.3")
    (format port "#version 130\n"))
   ((string>= version "1.2")
    (format port "#version 120\n"))
   (else
    (seagull-compiler-error "incompatible GLSL version"
                            (list version)
                            emit-version-preprocessor))))

(define (emit-shims version port)
  (when (string<= version "3.3")
    (format port "
vec4 texture(sampler2D tex, vec2 coord) {
  return texture2D(tex, coord);
}
vec4 texture(samplerCube tex, vec3 coord) {
  return textureCube(tex, coord);
}
")))

(define (emit-stage exp stage version)
  (call-with-output-string
    (lambda (port)
      (emit-version-preprocessor version port)
      (emit-shims version port)
      (emit-glsl exp 'fragment version port))))

(define* (link-seagull-modules vertex fragment version)
  (unless (seagull-module-vertex? vertex)
    (seagull-compiler-error "not a vertex shader"
                            (list vertex)
                            link-seagull-modules))
  (unless (seagull-module-fragment? fragment)
    (seagull-compiler-error "not a fragment shader"
                            (list fragment)
                            link-seagull-modules))
  (parameterize ((unique-identifier-counter
                  (max (seagull-module-max-id vertex)
                       (seagull-module-max-id fragment))))
    (unless (vertex-outputs-match-fragment-inputs? vertex fragment)
      (seagull-compiler-error "vertex outputs do not match fragment inputs"
                              (list vertex fragment)
                              link-seagull-modules))
    (unless (uniforms-compatible? vertex fragment)
      (seagull-compiler-error "vertex uniforms clash with fragment uniforms"
                              (list vertex fragment)
                              link-seagull-modules))
    (define-values (vertex* fragment* uniform-map)
      (link-vertex-outputs-with-fragment-inputs vertex fragment))
    (define vertex-glsl (emit-stage vertex* 'vertex version))
    (define fragment-glsl (emit-stage fragment* 'fragment version))
    (values vertex-glsl fragment-glsl uniform-map)))

(define* (compile-shader vertex fragment #:key
                         (version (graphics-engine-glsl-version)))
  (let-values (((glsl:vertex glsl:fragment uniform-map)
                (link-seagull-modules vertex fragment version)))
    (call-with-input-string glsl:vertex
      (lambda (vertex-port)
        (call-with-input-string glsl:fragment
          (lambda (fragment-port)
            (make-shader vertex-port fragment-port
                         #:uniform-map uniform-map
                         #:pre-process? #f)))))))


;;;
;;; REPL integration
;;;

(define-meta-command ((seagull-expand chickadee) repl stage exp)
  "seagull-expand STAGE EXP
Run the expander on EXP for shader STAGE."
  (parameterize ((unique-identifier-counter 0))
    (pretty-print (expand* exp stage))))

(define-meta-command ((seagull-simplify chickadee) repl stage exp)
  "seagull-simplify STAGE EXP
Run the partial evaluator on EXP for shader STAGE."
  (parameterize ((unique-identifier-counter 0))
    (pretty-print (simplify* (expand* exp stage)))))

(define-meta-command ((seagull-infer chickadee) repl stage exp)
  "seagull-infer STAGE EXP
Run type inference on EXP for shader STAGE."
  (parameterize ((unique-identifier-counter 0))
    (pretty-print
     (infer* (hoist-functions*
              (prune
               (simplify*
                (expand* exp stage))))
             stage))))

(define-meta-command ((seagull-inspect chickadee) repl module)
  "seagull-inspect MODULE
Show the intermediate compiled form of MODULE."
  (pretty-print (seagull-module-compiled (repl-eval repl module))))

(define-meta-command ((seagull-compile chickadee) repl vertex fragment
                      #:optional (version (graphics-engine-glsl-version)))
  "seagull-compile VERTEX-MODULE FRAGMENT-MODULE
Show the compiled GLSL form of VERTEX-MODULE and FRAGMENT-MODULE."
  (define-values (vertex-glsl fragment-glsl uniform-map)
    (link-seagull-modules (repl-eval repl vertex)
                          (repl-eval repl fragment)
                          version))
  (format #t "Vertex GLSL:\n\n~a\n" vertex-glsl)
  (format #t "Fragment GLSL:\n\n~a" fragment-glsl))