Add function hoisting and redo free variable analysis.

1 files changed, 120 insertions, 58 deletions

diff --git a/chickadee/graphics/seagull.scm b/chickadee/graphics/seagull.scm
index 72500fc..df5b9d1 100644
--- a/chickadee/graphics/seagull.scm
+++ b/chickadee/graphics/seagull.scm
@@ -35,6 +35,7 @@
   #:use-module (ice-9 format)
   #:use-module (ice-9 match)
   #:use-module (srfi srfi-1)
+  #:use-module (srfi srfi-11)
   #:export (compile-seagull))
 
 ;; The Seagull compiler is designed as a series of source-to-source
@@ -103,6 +104,21 @@
 (define (compose-envs . envs)
   (concatenate envs))
 
+(define (env-names env)
+  (map car env))
+
+(define (env-map proc env)
+  (map (match-lambda
+         ((name . exp)
+          (proc name exp)))
+       env))
+
+(define (env-for-each proc env)
+  (for-each (match-lambda
+              ((name . exp)
+               (proc name exp)))
+            env))
+
 (define (top-level-env)
   (empty-env))
 
@@ -218,28 +234,15 @@
        (make-exception-with-message "seagull: invalid syntax")
        (make-exception-with-irritants (list exp)))))))
 
-(define (expand* exp)
-  (parameterize ((unique-identifier-counter 0))
-    (expand exp (top-level-env))))
-
 
 ;;;
 ;;; Function hoisting
 ;;;
 
-;; Move all lambda bindings to the top-level.  As mentioned earlier,
-;; GLSL does not allow nested functions.
-
-
-;;;
-;;; Free variable annotation
-;;;
-
-;; All lambda forms are annotated with the variables that appear free
-;; in their body expressions.  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.
+;; 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)))
@@ -251,59 +254,112 @@
   (exception-accessor &seagull-scope-error
                       (record-accessor &seagull-scope-error 'variable)))
 
-(define (free-variables exp bound-vars top-level-vars)
+(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
-    ((? immediate?) '())
+    ((? immediate?)
+     #t)
     (('var name original-name)
-     (cond
-      ((memq name bound-vars) ; bound vars are not free
-       '())
-      ((memq name top-level-vars) ; free top-level var
-       (list name))
-      (else
-       ;; Free variables that aren't top-level are not allowed because
-       ;; GLSL doesn't support closures.
-       (raise-exception
-        (make-exception
-         (make-seagull-scope-error original-name)
-         (make-exception-with-origin free-variables)
-         (make-exception-with-message
-          "seagull: free variable is not top-level")
-         (make-exception-with-irritants (list exp)))))))
+     (or (memq name bound-vars) ; bound vars: OK
+         (memq name 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 original-name)
+           (make-exception-with-origin free-variables)
+           (make-exception-with-message
+            "seagull: free variable is not top-level")
+           (make-exception-with-irritants (list exp))))))
     (('if predicate consequent alternate)
-     (append (free-variables predicate bound-vars top-level-vars)
-             (free-variables consequent bound-vars top-level-vars)
-             (free-variables alternate bound-vars top-level-vars)))
-    (('let ((names _) ...) body)
-     (free-variables body (append names bound-vars) top-level-vars))
+     (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)))
     (('lambda (params ...) body)
-     (free-variables body (append params bound-vars) top-level-vars))
+     (check-free-variables body params top-level-vars))
     ((or ('primcall _ args ...)
          ('call args ...))
-     (append-map (lambda (arg)
-                   (free-variables arg bound-vars top-level-vars))
-                 args))))
-
-(define (annotate-free-variables exp)
+     (check-free-variables-in-list args 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))
+  (values `(let ,(map list names exps*) ,body*)
+          (compose-envs exps-env body-env)))
+
+(define (hoist-lambda params body)
+  (define var (unique-identifier))
+  (define-values (body* body-env)
+    (hoist-functions body))
+  (define lambda* `(lambda ,params ,body*))
+  (values `(var ,var #f)
+          (extend-env var lambda* body-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-functions exp)
   (match exp
-    ((or (? immediate?)
-         ('var _ _))
-     exp)
+    ((or (? immediate?) ('var _ _))
+     (values exp (empty-env)))
     (('if predicate consequent alternate)
-     `(if (annotate-free-variables predicate)
-          (annotate-free-variables consequent)
-          (annotate-free-variables alternate)))
+     (hoist-if predicate consequent alternate))
     (('let ((names exps) ...) body)
-     (define exps* (map annotate-free-variables exps))
-     `(let ,(map list names exps*) ,(annotate-free-variables body)))
+     (hoist-let names exps body))
     (('lambda (params ...) body)
-     ;; TODO: Actually figure out top-level vars.
-     (define free-vars (free-variables body params '(+ - * /)))
-     `(lambda ,free-vars ,params ,(annotate-free-variables body)))
+     (hoist-lambda params body))
     (('primcall operator args ...)
-     `(primcall ,operator ,@(map annotate-free-variables args)))
+     (hoist-primcall operator args))
     (('call args ...)
-     `(call ,@(map annotate-free-variables args)))))
+     (hoist-call args))))
+
+(define (hoist-functions* exp)
+  (define-values (exp* env)
+    (hoist-functions exp))
+  (define top-level-vars
+    (append (env-names env) '(+ - * /)))
+  (env-for-each (lambda (name exp)
+                  (check-free-variables exp '() top-level-vars))
+                env)
+  `(top-level ,(reverse (env-map list env))
+              ,exp*))
 
 
 ;;;
@@ -338,3 +394,9 @@
 
 ;; Combine all of the compiler passes on a user provided program and
 ;; emit GLSL code if the program is valid.
+
+(define (compile-seagull exp)
+  (parameterize ((unique-identifier-counter 0))
+    (let* ((expanded (expand exp (top-level-env)))
+           (hoisted (hoist-functions* expanded)))
+      hoisted)))