path: root/chickadee/graphics/seagull/pass-linearize.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 linearization pass transforms a tree of base IL objects into a
;; control flow graph of continuations.  Extremely based on Guile's
;; own "CPS soup" but greatly simplified for this much simpler
;; language and also my own inability to understand something more
;; complicated.
(define-module (chickadee graphics seagull pass-linearize)
  #:use-module (chickadee graphics seagull base)
  #:use-module (chickadee graphics seagull cps)
  #:use-module (chickadee graphics seagull types)
  #:use-module (chickadee graphics seagull utils)
  #:use-module (ice-9 match)
  #:use-module (language cps intmap)
  #:use-module (srfi srfi-1)
  #:use-module (srfi srfi-9)
  #:use-module (srfi srfi-11)
  #:export (linearize))

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

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

(define (next-var)
  (let ((var (var-counter)))
    (var-counter (+ var 1))
    var))

(define (fresh-var)
  (next-var))

(define (fresh-vars lst)
  (map (lambda (_x) (fresh-var)) lst))

(define (make-fresh-vars n)
  (list-tabulate n (lambda (_i) (fresh-var))))

;; Unique continuation label generation
(define label-counter (make-parameter 0))

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

(define (fresh-label)
  (let ((label (label-counter)))
    (label-counter (+ label 1))
    label))

(define (num-values exp)
  (match exp
    ((or (? assignment?)
         (? discard?))
     0)
    ((or (? constant?)
         (? lexical-reference?)
         (? struct-reference?)
         (? array-reference?)
         (? function?)
         (? primitive-call?)
         ($ <values> (_)))
     1)
    ;; Number of values is unknown at this stage. The type inference
    ;; pass will have to figure it out.
    (_ #f)))

(define (compile:constant src val graph k)
  (let* ((type (type-for-constant val))
         (const (make-cps-constant val type))
         (term (make-continue src k const)))
    (values graph term)))

(define (compile:conditional src predicate consequent alternate env graph k)
  (let ((branch-k (fresh-label))
        (consequent-k (fresh-label))
        (alternate-k (fresh-label)))
    (let*-values (((graph1 predicate-term)
                   (compile-cps predicate env graph branch-k))
                  ((graph2 consequent-term)
                   (compile-cps consequent env graph1 k))
                  ((graph3 alternate-term)
                   (compile-cps alternate env graph2 k)))
      (let* ((branch-var (fresh-var))
             (branch-term (make-branch src branch-var consequent-k alternate-k))
             (branch-cont (make-arguments (list branch-var) branch-term #f))
             (consequent-cont (make-arguments '() consequent-term #f))
             (alternate-cont (make-arguments '() alternate-term #f)))
        (values (intmap-add (intmap-add (intmap-add graph3 branch-k branch-cont)
                                        consequent-k consequent-cont)
                            alternate-k alternate-cont)
                predicate-term)))))

;; Converting a linear sequence of expressions to continuation passing
;; style requires an algorithm written in continuation passing style.
;; We have to traverse to the end of the list, and then chain together
;; the continuations as the algorithm works backwards.
(define (compile:arguments exps env graph return)
  (match exps
    (()
     (return '() graph))
    ((exp . rest)
     (let*-values (((name) (fresh-var))
                   ((graph* next-term)
                    (compile:arguments rest env graph
                                      (lambda (names graph)
                                        (return (cons name names) graph))))
                   ((k) (fresh-label))
                   ((cont) (make-arguments (list name) next-term #f))
                   ((graph** this-term)
                    (compile-cps exp env graph* k)))
       (values (intmap-add graph** k cont)
               this-term)))))

(define (compile:primitive-call src name args env graph k)
  (compile:arguments args env graph
                     (lambda (vars graph*)
                       (let* ((type (fresh-type-variable))
                              (prim (make-cps-primitive-call name vars type))
                              (term (make-continue src k prim)))
                         (values graph* term)))))

(define (compile:let src names exps body env graph k)
  (compile:arguments exps env graph
                     (lambda (vars graph*)
                       (let ((env* (fold extend env names vars)))
                         (compile-cps body env* graph* k)))))

(define (compile:lexical-reference src name env graph k)
  (let ((vals (make-cps-values (list (lookup name env)))))
    (values graph (make-continue src k vals))))

(define (compile:assignment src name value env graph k)
  (pk 'assignment name value)
  (let ((k* (fresh-label))
        (var (lookup name env)))
    (let*-values (((graph* term)
                   (compile-cps value env graph k*)))
      (let* ((val (fresh-var))
             (assign (make-cps-assignment var val))
             (term* (make-continue src k assign))
             (cont (make-arguments (list val) term* #f)))
        (values (intmap-add graph* k* cont) term)))))

(define (compile:sequence exps env graph k)
  ;; (match exps
  ;;   (()
  ;;    (return '() graph))
  ;;   ((exp . rest)
  ;;    (let*-values (((names) (and=> (num-values exp) make-fresh-vars))
  ;;                  ((graph* next-term)
  ;;                   (compile:arguments rest env graph
  ;;                                     (lambda (graph)
  ;;                                       (return graph))))
  ;;                  ((k) (fresh-label))
  ;;                  ((cont) (make-arguments (list name) next-term #f))
  ;;                  ((graph** this-term)
  ;;                   (compile-cps exp env graph* k)))
  ;;      (values (intmap-add graph** k cont)
  ;;              this-term))))

  (match exps
    ((exp)
     (pk 'seq exp)
     (compile-cps exp env graph k))
    ((exp . rest)
     (pk 'seq exp)
     (let ((vars (pk 'seq-vars (and=> (pk 'num-values exp (num-values exp)) make-fresh-vars))))
       (let-values (((graph* term)
                     (compile:sequence rest env graph k)))
         (let* ((k* (fresh-label))
                (cont (make-arguments vars term #f)))
           (compile-cps exp env (intmap-add graph* k* cont) k*)))))))

'((0 . (arguments () (continue 4 (constant 1))))
  (1 . (arguments (2) (return 2)))
  (2 . (arguments () (continue 1 (values 0))))
  (3 . (arguments (1) (continue 2 (assignment 0 1))))
  (4 . (arguments (0) (continue 3 (constant 2)))))

(define (compile-cps exp env graph k)
  (match (pk 'compile exp)
    (($ <constant> src val)
     (compile:constant src val graph k))
    (($ <conditional> src predicate consequent alternate)
     (compile:conditional src predicate consequent alternate env graph k))
    (($ <primitive-call> src name args)
     (compile:primitive-call src name args env graph k))
    (($ <sequence> src exps)
     (compile:sequence exps env graph k))
    (($ <let> src _ names exps body)
     (compile:let src names exps body env graph k))
    (($ <lexical-reference> src _ name)
     (compile:lexical-reference src name env graph k))
    (($ <assignment> src _ name value)
     (compile:assignment src name value env graph k))))

;; (define (term-arity term)
;;   (define (exp-arity exp)
;;     (match exp
;;       (($ <cps-constant> _ _)
;;        1)
;;       (($ <cps-assignment> _ _)
;;        0)))
;;   (match term
;;     (($ <continue> _ _ exp)
;;      (exp-arity exp))
;;     (($ ))))

(define-syntax intmap-add*
  (syntax-rules ()
    ((_ intmap) intmap)
    ((_ intmap (i v) rest ...)
     (intmap-add (intmap-add* intmap rest ...) i v))))

(define (linearize exp)
  (parameterize ((var-counter 0)
                 (label-counter 0))
    ;; The initial continuation is always at index 0.
    (let* ((init-label (fresh-label))
           (func-return-label (fresh-label)))
      (let-values (((graph term)
                    (compile-cps exp (fresh-environment) empty-intmap func-return-label)))
        (let* ((func-start-label (fresh-label))
               (func-label (fresh-label))
               (end-label (fresh-label))
               (params '())
               (results (list (fresh-var)))
               (init-term (make-continue #f end-label (make-cps-function func-label)))
               (init-cont (make-arguments '() init-term #f))
               (func-cont (make-function-entry #f params results func-start-label func-return-label #f))
               (func-start-cont (make-arguments '() term #f))
               (func-return-term (make-return results))
               (func-return-cont (make-arguments results func-return-term #f))
               (end-results (list (fresh-var)))
               (end-term (make-return '()))
               (end-cont (make-arguments end-results end-term #f)))
          (intmap-add* graph
                       (init-label init-cont)
                       (func-label func-cont)
                       (func-start-label func-start-cont)
                       (func-return-label func-return-cont)
                       (end-label end-cont)))))))