path: root/chickadee/graphics/seagull/pass-expand.scm

;;; Chickadee Game Toolkit
;;; Copyright © 2023 David Thompson <dthompson2@worcester.edu>
;;;
;;; Licensed under the Apache License, Version 2.0 (the "License");
;;; you may not use this file except in compliance with the License.
;;; You may obtain a copy of the License at
;;;
;;;    http://www.apache.org/licenses/LICENSE-2.0
;;;
;;; Unless required by applicable law or agreed to in writing, software
;;; distributed under the License is distributed on an "AS IS" BASIS,
;;; WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
;;; See the License for the specific language governing permissions and
;;; limitations under the License.

;; The expansion pass transforms a tree of Seagull syntax objects into
;; a tree of primitives in the format of the base intermediate
;; language.  The expander also takes care of "alpha conversion",
;; which is a fancy math way of saying that all names are given new,
;; unique names so that the compiler doesn't have to worry about name
;; collisions, like if '+' refers to the primitive addition function
;; or a lexical variable that shadows it.
(define-module (chickadee graphics seagull pass-expand)
  #:use-module (chickadee graphics seagull base)
  #:use-module (chickadee graphics seagull primitives)
  #:use-module (chickadee graphics seagull syntax)
  #:use-module (chickadee graphics seagull utils)
  #:use-module (ice-9 match)
  #:use-module (srfi srfi-1)
  #:export (expand))

;; Unique name generation.
(define name-counter (make-parameter 0))

(define (reset-name-counter!)
  (name-counter 0))

(define (next-name-id)
  (let ((id (name-counter)))
    (name-counter (+ id 1))
    id))

(define (fresh-name)
  (string->symbol (format #f "t~a" (next-name-id))))

(define (fresh-names names)
  (map (lambda (_name) (fresh-name)) names))

(define (alpha-convert names)
  (define names* (map (lambda (_name) (fresh-name)) names))
  (fold extend (fresh-environment) names names*))

(define (expand:list exps stage env)
  (map (lambda (x) (expand* x stage env)) exps))

(define (expand:body body src stage env)
  (match body
    ;; Empty body, invalid.
    (()
     (seagull-syntax-error "body requires at least one expression"
                           (build-syntax src '())
                           expand:discard))
    ;; Body with a single expression, just expand that expression.
    ((exp)
     (expand* exp stage env))
    ;; Body with multiple expressions, expand to a sequence.
    ((exps ...)
     (make-sequence src (expand:list exps stage env)))))

;; Lookup alpha-converted name for the given source name.
(define (expand:reference name src stage env)
  (let ((name* (lookup name env)))
    (unless name*
      (seagull-syntax-error "unbound variable"
                            (make-seagull-syntax src name)
                            expand:reference))
    (make-lexical-reference src name name*)))

(define (expand:if predicate consequent alternate src stage env)
  (make-conditional src
                    (expand* predicate stage env)
                    (expand* consequent stage env)
                    (expand* alternate stage env)))

(define (expand:function params body src stage env)
  (define params* (fresh-names params))
  (define env* (fold extend env params params*))
  (make-function src params params*
                 (expand:body body src stage env*)))

(define (expand:primcall operator arguments src stage env)
  (let* ((op (lookup-primitive-operator operator))
         (op:expand (primitive-operator-expand op)))
    (op:expand arguments src (lambda (stx) (expand* stx stage env)))))

(define (expand:call function arguments src stage env)
  (make-call src (expand* function stage env)
             (expand:list arguments stage env)))

(define (expand:let names values body src stage env)
  (match names
    (()
     (expand:body body src stage env))
    (_
     (let* ((values* (expand:list values stage env))
            (names* (fresh-names names))
            (env* (fold extend env names names*)))
       (make-let src names names* values* (expand:body body src stage env*))))))

(define (expand:let-values names values body src stage env)
  (match names
    (()
     (expand:body body src stage env))
    (_
     (let* ((values* (expand:list values stage env))
            (names* (map fresh-names names))
            (env* (fold extend env (concatenate names) (concatenate names*))))
       (make-let-values src names names* values* (expand:body body src stage env*))))))

(define (expand:values exps src stage env)
  (make-values src (expand:list exps stage env)))

(define (expand:struct-ref value fields src stage env)
  (define value* (expand* value stage env))
  (match fields
    ((field . rest)
     (let loop ((fields rest)
                (prev (make-struct-reference src value* field)))
       (match fields
         (() prev)
         ((field . rest)
          (loop rest (make-struct-reference src prev field))))))))

(define (expand:array-ref value indices src stage env)
  (define value* (expand* value stage env))
  (match indices
    ((i . rest)
     (let ((i* (expand* i stage env)))
       (let loop ((indices rest)
                  (prev (make-array-reference src value* i*)))
         (match indices
           (() prev)
           ((j . rest)
            (let ((j* (expand* j stage env)))
              (loop rest (make-array-reference src prev j*))))))))))

(define (expand:assignment name value src stage env)
  (make-assignment src name (lookup name env) (expand* value stage env)))

(define (expand:discard src stage)
  (unless (fragment-stage? stage)
    (seagull-syntax-error "'discard' is only permitted in fragment shaders"
                          (build-syntax src '(discard))
                          expand:discard)))

;; Macros:
(define (expand:let* bindings body src stage env)
  (match bindings
    (()
     (expand* body stage env))
    ((binding . rest)
     (expand* (build-syntax (seagull-syntax-source binding)
                            `(let (,binding)
                               (let* ,rest ,body)))
              stage env))))

(define (expand:or exps src stage env)
  (match exps
    (()
     (make-constant src #f))
    ((exp)
     (expand* exp stage env))
    ((exp . rest)
     (expand* (build-syntax src `(let ((x ,exp)) (if x x (or ,@rest))))
              stage env))))

(define (expand:and exps src stage env)
  (match exps
    (()
     (make-constant src #t))
    ((exp)
     (expand* exp stage env))
    ((exp . rest)
     (expand* (build-syntax src `(if ,exp (and ,@rest) #f))
              stage env))))

(define (expand:cond clauses src 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.
      ((($ <syntax> _ (($ <syntax> _ 'else) exp)))
       exp)
      ((($ <syntax> _ (predicate consequent)) . rest)
       (build-syntax src `(if ,predicate ,consequent ,(cond->if rest))))
      (()
       (seagull-syntax-error "'cond' form must end with 'else' clause"
                             (build-syntax src `(cond ,@clauses))
                             expand:cond))
      (_
       (seagull-syntax-error "invalid 'cond' form"
                             (build-syntax src `(cond ,@clauses))
                             expand:cond))))
  (expand* (cond->if clauses) stage env))

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

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

(define (primitive-operator-for-stage? name stage)
  (let ((op (lookup-primitive-operator name)))
    (and op (memq stage (primitive-operator-stages op)))))

(define (expand* syntax stage env)
  (define (lexically-bound? x)
    (and (symbol? x) (bound? x env)))
  (define (primop? x)
    (and (symbol? x) (primitive-operator-for-stage? x stage)))
  (define src (seagull-syntax-source syntax))
  (match (seagull-syntax-expression syntax)
    ;; Constants and variable references
    ((? primitive-constant? c)
     (make-constant src c))
    ((? symbol? name)
     (expand:reference name src stage env))
    ;; Function call with a variable reference in operator position
    ;; that *might* shadow a built-in, so we need to check for it
    ;; early.
    (((and ($ <syntax> name-src (? lexically-bound?)) function) arguments ...)
     (expand:call function arguments src stage env))
    ;; Primitive syntax
    ((($ <syntax> _ 'if) predicate consequent alternate)
     (expand:if predicate consequent alternate src stage env))
    ((($ <syntax> _ 'lambda)
      ($ <syntax> _ (($ <syntax> _ (? symbol? params)) ...))
      body ...)
     (expand:function params body src stage env))
    ((($ <syntax> _ 'let)
      ($ <syntax> _ (($ <syntax> _ (($ <syntax> _ (? symbol? names))
                                    values))
                     ...))
      body ...)
     (expand:let names values body src stage env))
    ((($ <syntax> _ 'let-values)
      ($ <syntax> _ (($ <syntax> _ (($ <syntax> _ (($ <syntax> _ (? symbol? names)) ...))
                                    values))
                     ...))
      body ...)
     (expand:let-values names values body src stage env))
    ((($ <syntax> _ 'values) exps ...)
     (expand:values exps src stage env))
    ((($ <syntax> _ '->) value ($ <syntax> _ (? symbol? fields)) ..1)
     (expand:struct-ref value fields src stage env))
    ((($ <syntax> _ '@) value indices ..1)
     (expand:array-ref value indices src stage env))
    ((($ <syntax> _ 'set!) ($ <syntax> _ (? symbol? name)) value)
     (expand:assignment name value src stage env))
    ((($ <syntax> _ 'discard))
     (expand:discard src stage))
    ;; Built-in macros
    ((($ <syntax> _ 'let*) ($ <syntax> _ (bindings ...)) body)
     (expand:let* bindings body src stage env))
    ((($ <syntax> _ 'and) exps ...)
     (expand:and exps src stage env))
    ((($ <syntax> _ 'or) exps ...)
     (expand:or exps src stage env))
    ((($ <syntax> _ 'cond) clauses ...)
     (expand:cond clauses src stage env))
    ((($ <syntax> _ 'case) key clauses ...)
     (expand:case key clauses src stage env))
    ;; Function calls
    ((($ <syntax> _ (? primop? op)) arguments ...)
     (expand:primcall op arguments src stage env))
    ((function arguments ...)
     (expand:call function arguments src stage env))
    ;; Uh oh
    (_
     (seagull-syntax-error "unknown form" syntax expand*))))

(define (expand syntax stage)
  (expand* syntax stage (fresh-environment)))