Add array references.

1 files changed, 134 insertions, 14 deletions

diff --git a/chickadee/graphics/seagull.scm b/chickadee/graphics/seagull.scm
index 8e3c613..ce1a686 100644
--- a/chickadee/graphics/seagull.scm
+++ b/chickadee/graphics/seagull.scm
@@ -302,6 +302,18 @@
           (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)))))))
+
 ;; Arithmetic operators, in true Lisp fashion, can accept many
 ;; arguments.  + and * accept 0 or more. - and / accept one or more.
 ;; The expansion pass transforms all such expressions into binary
@@ -449,6 +461,8 @@
     ;; Macros:
     (('-> exp (? symbol? members) ..1)
      (expand:-> exp members stage env))
+    (('@ exp indices ...)
+     (expand:@ exp indices stage env))
     (('let* (bindings ...) body)
      (expand:let* bindings body stage env))
     (('+ args ...)
@@ -541,6 +555,10 @@
 (define (propagate:struct-ref exp field env)
   `(struct-ref ,(propagate-constants exp env) ,field))
 
+(define (propagate:array-ref array-exp index-exp env)
+  `(array-ref ,(propagate-constants array-exp env)
+              ,(propagate-constants index-exp env)))
+
 ;; 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.
@@ -593,6 +611,8 @@
      (propagate:call operator args env))
     (('struct-ref exp field)
      (propagate:struct-ref exp field env))
+    (('array-ref array-exp index-exp)
+     (propagate:array-ref array-exp index-exp env))
     (('outputs (names exps) ...)
      (propagate:outputs names exps env))
     (('top-level inputs body)
@@ -656,6 +676,9 @@
      (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)
@@ -733,6 +756,14 @@
   (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))
@@ -766,6 +797,8 @@
      (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)
@@ -846,6 +879,23 @@
 (define-syntax-rule (define-struct-type (var-name name) (types names) ...)
   (define var-name (struct-type 'name (list (cons 'names types) ...))))
 
+;; Array type:
+(define (array-type type length)
+  `(array ,type ,length))
+
+(define (array-type? type)
+  (match type
+    (('array _ _) #t)
+    (_ #f)))
+
+(define (array-type-ref type)
+  (match type
+    (('array type _) type)))
+
+(define (array-type-length type)
+  (match type
+    (('array _ n) n)))
+
 ;; Type variables:
 (define unique-variable-type-counter (make-parameter 0))
 
@@ -946,6 +996,9 @@
      (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 (error "invalid type" type))))
@@ -1112,6 +1165,9 @@
 (define (predicate:struct-has? struct field var)
   `(struct-has? ,struct ,field ,var))
 
+(define (predicate:array-element array var)
+  `(array-element ,array ,var))
+
 (define (compose-predicates a b)
   (cond
    ((and (eq? a #t) (eq? b #t))
@@ -1154,7 +1210,10 @@
     (('struct-has? struct field var)
      `(struct-has? ,(apply-substitution-to-type struct from to)
                    ,field
-                   ,var))))
+                   ,(apply-substitution-to-type var from to)))
+    (('array-element array var)
+     `(array-element ,(apply-substitution-to-type array from to)
+                     ,(apply-substitution-to-type var from to)))))
 
 (define (apply-substitutions-to-predicate pred subs)
   (env-fold (lambda (from to pred*)
@@ -1321,6 +1380,12 @@
     (('struct-has? struct field field-var)
      (if (struct-type? struct)
          (values #t (list (cons field-var (struct-type-ref struct field))))
+         (values pred '())))
+    ;; Substitute the element var when array has been resolved to an
+    ;; array type.
+    (('array-element array element-var)
+     (if (array-type? array)
+         (values #t (list (cons element-var (array-type-ref array))))
          (values pred '())))))
 
 (define (eval-predicate* pred subs)
@@ -1342,6 +1407,8 @@
    ((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)))
@@ -1404,6 +1471,9 @@
              (free-variables-in-type b)))
     (('struct-has? 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
@@ -1684,6 +1754,34 @@
           (compose-predicates exp-pred
                               (predicate:struct-has? exp-type field tvar))))
 
+(define (infer:array-ref array-exp index-exp env)
+  (define-values (array-exp* array-exp-subs array-exp-pred)
+    (infer-exp array-exp env))
+  (define array-type (single-type array-exp*))
+  (define env* (apply-substitutions-to-env env array-exp-subs))
+  (define-values (index-exp* index-exp-subs index-exp-pred)
+    (infer-exp index-exp env*))
+  (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* (compose-predicates (predicate:array-element array-type tvar)
+                                         (compose-predicates 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))
@@ -1747,8 +1845,8 @@
                        combined-subs
                        new-pred
                        env*))
-      (((_ type-name name) . rest)
-       (define types (list (type-name->type type-name)))
+      (((_ desc name) . rest)
+       (define types (list (type-descriptor->type desc)))
        (infer-bindings rest
                        (cons types texps)
                        subs
@@ -1808,6 +1906,8 @@
      (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)
@@ -1836,17 +1936,19 @@
 (define type:mat4 (primitive-type 'mat4))
 (define type:sampler-2d (primitive-type 'sampler-2d))
 
-(define (type-name->type name)
-  (case name
-    ((bool) type:bool)
-    ((int) type:int)
-    ((float) type:float)
-    ((vec2) type:vec2)
-    ((vec3) type:vec3)
-    ((vec4) type:vec4)
-    ((mat3) type:mat3)
-    ((mat4) type:mat4)
-    ((sampler-2d) type:sampler-2d)))
+(define (type-descriptor->type desc)
+  (match desc
+    ('bool type:bool)
+    ('int type:int)
+    ('float type:float)
+    ('vec2 type:vec2)
+    ('vec3 type:vec3)
+    ('vec4 type:vec4)
+    ('mat3 type:mat3)
+    ('mat4 type:mat4)
+    ('sampler-2d type:sampler-2d)
+    (('array desc* (? exact-integer? length))
+     (array-type (type-descriptor->type desc*) length))))
 
 (define-syntax-rule (a+b->c (ta tb tc) ...)
   (let ((a (fresh-variable-type))
@@ -2300,6 +2402,22 @@
           field)
   (list output-temp))
 
+(define (emit:array-ref type array-exp index-exp version port level)
+  (define array-temp
+    (match (emit-glsl array-exp version port level)
+      ((temp) temp)))
+  (define index-temp
+    (match (emit-glsl index-exp version port level)
+      ((temp) temp)))
+  (define output-temp (unique-identifier))
+  (indent level port)
+  (format port "~a ~a = ~a[~a];\n"
+          (type-name type)
+          output-temp
+          array-temp
+          index-temp)
+  (list output-temp))
+
 (define %type-name-map
   '((sampler-2d . sampler2D)))
 
@@ -2360,6 +2478,8 @@
      (emit:call types operator args version port level))
     (('t (type) ('struct-ref exp field))
      (emit:struct-ref type exp field version port level))
+    (('t (type) ('array-ref array-exp index-exp))
+     (emit:array-ref type array-exp index-exp version port level))
     (('t _ ('outputs (names exps) ...))
      (emit:outputs names exps version port level))
     (('t _ ('top-level (bindings ...) body))