;;;; Copyright 2002, Jonathan Bachrach.  See file TERMS.
;;;; BASED ON Lisp In Small Pieces by Christian Queinnec

(use goo)
(use goo/boot)
(use eval/ast)

;;; BOXING

(dc <box-read> (<computed-program>))
  (dp! box-reference (x|<box-read> => <any>))

(dc <box-write> (<computed-program>))
  (dp! box-reference (x|<box-write> => <any>))
  (dp! box-form (x|<box-write> => <any>))

(dc <box-creation> (<computed-program>))
  (dp! box-reference (x|<box-creation> => <any>))

(dm insert-box! (o|<program>)
  (ast-walk! insert-box! o))

(dm insert-box! (o|<local-reference>)
  (if (binding-mutable? (reference-binding o))
      (new <box-read> box-reference o)
      o))

;;; IDENTIFY LOCAL CALLS
(dm insert-box! (o|<regular-application>)
  (ast-walk! insert-box! o)
  (when (isa? (application-function o) <local-reference>)
    (set (reference-called-function? (application-function o)) #t))
  o)

(dm insert-box! (o|<local-assignment>)
  (new <box-write> 
    box-reference (assignment-reference o)
    box-form      (insert-box! (assignment-form o))))

(dm insert-box! (o|<ast-method>)
  (set (function-body o)
       (insert-box! (boxify-mutable-bindings
                     (function-body o) (function-bindings o))))
  o)

(dm insert-box! (o|<fix-let>)
  (set (fix-let-arguments o)
       (insert-box! (fix-let-arguments o)))
  (set (fix-let-body o)
       (insert-box! 
        (boxify-mutable-bindings (fix-let-body o) (fix-let-bindings o))))
  o)

(dm insert-box! (o|<locals>)
  (set (locals-functions o)
       (map insert-box! (locals-functions o)))
  (set (locals-body o)
       (insert-box!
        (boxify-mutable-bindings (locals-body o) (locals-bindings o))))
  o)

(dm boxify-mutable-bindings (form bindings|<lst>)
  (sequentialize
    (rep loop ((res '()) (bindings bindings))
      (if (nul? bindings)
          (rev! (pair form res))
          (let ((binding (head bindings)))
            (if (binding-mutable? binding)
                (let ((creator 
                        (new <box-creation>
                          box-reference
                           (new <local-reference> reference-binding binding))))
                  (loop (pair creator res) (tail bindings)))
                (loop res (tail bindings))))))))

;;; FLATTENING ENVIRONMENTS

(dc <free-reference> (<real-reference>))
  (dp! reference-offset (x|<free-reference> => (t? <int>))) 
  (dp! reference-self? (x|<free-reference> => <log>))

(dm lift! (o)
  (lift-procedures! o #f '()))

(dm lift-procedures! (o|<program> flat-fun bindings)
  (ast-walk! lift-procedures! o flat-fun bindings))

(dm lift-procedures! (o|<local-reference> flat-fun bindings)
  (def b (reference-binding o))
  (if (mem? bindings b)
      o
      (let ((offset (adjoin-free-binding! flat-fun o)))
        (set (binding-closed-over? b) #t)
        (new <free-reference> 
          reference-binding b
          reference-offset  offset
          reference-self?
            (== (reference-binding o) (function-binding flat-fun))))))

(dm adjoin-free-binding! (flat-fun|<ast-method> ref => <int>)
  (rep check ((i|<int> 0) (free* (function-free flat-fun)))
    (if (nul? free*)
        (let ((new-env (free-environment ref (nul <free-environment>))))
          (if (nul? (function-free flat-fun))
              (set (function-free flat-fun) new-env)
              (rep add ((free* (function-free flat-fun)))
                (def tail (tail free*))
                (if (nul? tail)
                    (set (tail free*) new-env)
                    (add tail))))
          i)
        (if (== (reference-binding ref)
                (reference-binding (head free*)))
            i
            (check (+ i 1) (tail free*))))))

(dm lift-procedures! (o|<fix-let> flat-fun bindings)
  (set (fix-let-arguments o)
       (lift-procedures! (fix-let-arguments o) flat-fun bindings))
  (def new-bindings (cat (fix-let-bindings o) bindings))
  (set (fix-let-body o)
       (lift-procedures! (fix-let-body o) flat-fun new-bindings))
  o)

(dm lift-procedures! (o|<locals> flat-fun bindings)
  (def new-bindings (cat (locals-bindings o) bindings))
  (set (locals-functions o)
       (map (rcurry lift-procedures! flat-fun new-bindings) 
            (locals-functions o)))
  (set (locals-body o)
       (lift-procedures! (locals-body o) flat-fun new-bindings))
  o)

(dm lift-procedures! (o|<ast-primitive> flat-fun bindings)
  (set (function-body o)
       (lift-procedures! (function-body o) #f (function-bindings o)))
  o)

(dm lift-procedures! (o|<ast-generic> flat-fun bindings)
  o)

(dm lift-procedures! (o|<ast-method> flat-fun bindings)
  (set (function-body o)
       (lift-procedures! (function-body o) o (function-bindings o)))
  (set (function-free o) 
       (lift-procedures! (function-free o) flat-fun bindings))
  o)

;;; COLLECTING QUOTATIONS AND FUNCTIONS

(dc <flattened-program> (<computed-program>))
  (dp! program-form (x|<flattened-program> => (t? <program>)))
  (dp! program-quotations (x|<flattened-program> => <any>) (fab <tab> 40))
  (dp! program-bindings (x|<flattened-program> => <any>) (fab <tab> 40))
  (dp! program-definitions (x|<flattened-program> => <any>) '())

(dc <top-level-form> (<computed-program>))
  (dp! form-program (x|<top-level-form> => <program>))
  (dp! form-quotations (x|<top-level-form> => <any>) '())
  (dp! form-definitions (x|<top-level-form> => <any>) '())

(dc <closure-creation> (<computed-program>))
  (dp! closure-creation-index (x|<closure-creation> => <any>))
  (dp! closure-creation-free (x|<closure-creation> => <any>))

(dm flatten-seqs (o)
  (lst o))

(dm flatten-seqs (o|<sequential>)
  (rep loop ((s o) (r '()))
     (if (nul? s)
         (rev! r)
         (rep inner ((ss (flatten-seqs (head s))) (r r))
           (if (nul? ss)
               (loop (tail s) r)
               (inner (tail ss) (pair (head ss) r)))))))

(dm as-top-level-forms (programs)
  (map (fun (x) (new <top-level-form> form-program x)) programs))

(dm extract-things! (o)
  (def forms  (as-top-level-forms (flatten-seqs o)))
  (def result (new <flattened-program>))
  (set (program-form result)
       (sequentialize (map (fun (form) (extract! form form result)) forms)))
  result)

(dm extract! (o|<program> form|<top-level-form> result|<flattened-program>)
  (ast-walk! extract! o form result))

(dm extract! (o|<constant> form|<top-level-form> result|<flattened-program>)
  (def value (constant-value o))
  (if (or (== value #f) (== value #t) (== value '())
          (isa? value <int>) (isa? value <chr>))
      o
      (let ((qb*   (program-quotations result))
            (index (len qb*))
            (probe (elt-or qb* value #f)))
        (if probe
            (new <global-reference> reference-binding probe)
            (let ((qb (new <module-binding> 
                        binding-kind 'quotation
                        binding-name index 
                        binding-info value)))
              (set (elt qb* value) qb)
              (pushf (form-quotations form) qb)
              (new <global-reference> reference-binding qb))))))

(dm extract! (o|<raw-constant> form|<top-level-form> result|<flattened-program>)
  o)

(dm extract! (o|<ast-method> form|<top-level-form> result|<flattened-program>)
  (set (function-debug-name o)
       (extract! (function-debug-name o) form result))
  (set (function-signature o)
       (extract! (function-signature o) form result))
  (set (function-body o) (extract! (function-body o) form result))
  (do (fun (binding)
        (set (binding-type binding) 
             (extract! (binding-type binding) form result)))
      (function-bindings o))
  (def free          (function-free o))
  (def free-bindings (map reference-binding free))
  (def index         (adjoin-definition! form result o))
  (set (function-index o) index)
  (set (function-free o)  free-bindings)
  (new <closure-creation> 
    closure-creation-index index closure-creation-free free))

(dm extract! 
    (o|<ast-primitive-definition> 
     form|<top-level-form> result|<flattened-program>)
  (def prim (assignment-form o))
  (set (function-body prim) (extract! (function-body prim) form result))
  (pushf (program-definitions result) prim)
  (extract! (new <constant> constant-value #f) form result))

(dm adjoin-definition! (form|<top-level-form> result|<flattened-program> x)
  (def new-index (len (program-definitions result)))
  (pushf (program-definitions result) x)
  (pushf (form-definitions form)      x)
  new-index)

;;; CLOSURIZE MAIN

(dm split-program (o max-count|<int>)
  (lst o))

(dm split-program (o|<sequential> max-count|<int>)
  (rep loop ((forms '()) (s o))
    (rep inner ((res '()) (es s) (count 0))
      (cond ((nul? es)
             (rev! (pair (sequentialize (rev! res)) forms)))
            ((= count max-count)
             (loop (pair (sequentialize (rev! res)) forms) es))
            (#T 
              (inner (pair (head es) res) (tail es) (+ count 1)))))))

(dm closurize-main! (o|<flattened-program> r max-count => <flattened-program>)
  (def forms      (split-program (program-form o) max-count))
  (def base-index (len (program-definitions o))) ;; TODO: HACK INTERFACE
  (rep loop ((calls '()) (defns (program-definitions o)) (i 0) (forms forms))
    (if (nul? forms)
        (seq (set (program-definitions o) (rev! defns))
             (set (program-form o)
                  (new <top-level-form>
                    form-program     (sequentialize (rev! calls))
                    form-definitions (rev! defns)))
             o)
        (let ((name
               (cat-sym "---main-" i "---"))
              (binding
               (ast-define-binding r name #f 'predefined))
              (defn
               (new <ast-primitive>
                 function-binding    binding
                 function-debug-name (objectify-quotation name 4)
                 function-signature  (objectify-signature '() r)
                 function-body       (head forms)))
              (call
               (new <predefined-application> 
                 application-binding   binding
                 application-arguments (nul <arguments>)
                 application-tail?     #f)))
          (loop (pair call calls) (pair defn defns) (+ i 1) (tail forms))))))

;; (def f (fun (n k) (if (= n 0) (k 1) (f (- n 1) (fun (r) (k (* n r)))))))

;;; GATHER-TEMPORARIES -- PULLS ALL FIX-LET ARGS UP TO ENCLOSING FUN
;;;   FLAT-FUN IS ENCLOSING FUN AND R IS ENVIRONMENT FOR NEW AST

(dc <assoc> (<any>))
  (dp! assoc-key (x|<assoc> => <any>))
  (dp! assoc-value (x|<assoc> => <any>))
(df assoc (x y => <assoc>)
  (new <assoc> assoc-key x assoc-value y))

(dm assocq (x l|<lst> => <any>)
  (if (nul? l)
      #f
      (if (== x (assoc-key (head l)))
          (head l)
          (assocq x (tail l)))))

(dm gather-temporaries! (o|<program>)
  (set (program-definitions o)
       (map (fun (def) (collect-temporaries! def def '()))
            (program-definitions o)))
  o)

(dm collect-temporaries! (o|<program> flat-fun r)
  (ast-walk! collect-temporaries! o flat-fun r))

(dm collect-temporaries! (o|<local-reference> flat-fun r)
  (def binding (reference-binding o))
  (def b       (assocq binding r))
  (if b
      (new <local-reference> reference-binding (assoc-value b))
      o))

(df prepend-renamings (x y r)
  ;; (cat (map2 assoc x y) r) ;; speed critical
  (rep loop ((z r) (x x) (y y))
    (if (nul? x)
        z
        (loop (pair (assoc (head x) (head y)) z) (tail x) (tail y)))))

(dm collect-temporaries! (o|<fix-let> flat-fun r)
  (set (fix-let-arguments o)
       (collect-temporaries! (fix-let-arguments o) flat-fun r))
  (def new-bindings (map new-renamed-binding (fix-let-bindings o)))
  (def new-r (prepend-renamings (fix-let-bindings o) new-bindings r))
  (adjoin-temporary-variables! flat-fun new-bindings)
  (set (fix-let-bindings o) new-bindings)
  (set (fix-let-body o) (collect-temporaries! (fix-let-body o) flat-fun new-r))
  o)

(dm collect-temporaries! (o|<loop> flat-fun r)
  (set (loop-bindings o)
       (map (fun (b) 
              (def nb (assocq b r))
              (if nb (assoc-value nb) b))
            (loop-bindings o)))
  (set (loop-body o) (collect-temporaries! (loop-body o) flat-fun r))
  o)

(dm collect-temporaries! (o|<locals> flat-fun r)
  (def new-bindings (map new-renamed-binding (locals-bindings o)))
  (def new-r (prepend-renamings (locals-bindings o) new-bindings r))
  (set (locals-functions o)
       (map (rcurry collect-temporaries! flat-fun new-r) 
            (locals-functions o)))
  (adjoin-temporary-variables! flat-fun new-bindings)
  (set (locals-bindings o) new-bindings)
  (set (locals-body o)
       (collect-temporaries! (locals-body o) flat-fun new-r))
  o)

(dm adjoin-temporary-variables! (flat-fun new-bindings)
  (rep adjoin 
      ((temps (function-temporaries flat-fun)) (bindings new-bindings))
    (if (nul? bindings)
        (set (function-temporaries flat-fun) temps)
        (if (mem? temps (head bindings))
            (adjoin temps (tail bindings))
            (adjoin (pair (head bindings) temps) (tail bindings))))))

(dc <renamed-local-binding> (<local-binding>))
  (dp! binding-index (x|<renamed-local-binding> => <any>))

(ddv *renaming-bindings-counter* 0)

(dm new-renamed-binding (binding|<local-binding>)
  (set *renaming-bindings-counter* (+ *renaming-bindings-counter* 1))
  (new <renamed-local-binding>      ;; TODO: SHOULD USE CLONING
    binding-name            (binding-name binding)
    binding-type            (binding-type binding)
    binding-dynamic-extent? (binding-dynamic-extent? binding)
    binding-closed-over?    (binding-closed-over? binding)
    binding-index           *renaming-bindings-counter*))

;;; REGISTER ALLOCATION

(dv *register-passive?* #f) ;; FORCES ALL LOADS THROUGH REGISTERS (E.G., RISC)

(dm register-allocate! (o|<program> register-passive?|<log>)
  (dlet ((*register-passive?* register-passive?))
    (do (fun (def) (collect-registers! def #f #t)) (program-definitions o)))
  o)

(df allocate-register (o|<program> f val?)
  (when (and val? f (not (program-register o)))
    ;; TODO: INEFFICIENT
    (set (program-register o) (len (function-registers f)))
    (add! (function-registers f) o)))

(dm collect-registers! (o|<computed-program> f val?)
  (allocate-register o f val?)
  (ast-walk collect-registers! o f #t))

(dm collect-registers! (o|<sequential> f val?)
  (rep loop ((x o))
    (if (nul? x)
        o
        (let ((next-x (tail x))) ;; ONLY LAST FORM IN SEQ CONSUMES VALUE
          (collect-registers! (head x) f (and val? (nul? next-x)))
          (loop next-x)))))

(dm collect-registers! (o|<passive-program> f val?)
  (when *register-passive?* 
    (allocate-register o f val?))
  (ast-walk collect-registers! o f #t))

(dm collect-registers! (o|<top-level-form> f val?)
  (allocate-register o f val?)
  (do (fun (def) (collect-registers! def f val?)) (form-definitions o))
  (collect-registers! (form-program o) f val?))

(dm collect-registers! (o|<ast-method> f val?)
  (allocate-register o f val?)
  (collect-registers! (function-signature o) f #t)
  (collect-registers! (function-debug-name o) o #t)
  (collect-registers! (function-body o) o #t))

(dm collect-registers! (o|<ast-primitive> f val?)
  (collect-registers! (function-body o) o #t))

(dm collect-registers! (o|<locals> f val?)
  (collect-registers! (locals-body o) f #t)
  (collect-registers! (locals-functions o) f #t))

(dm collect-registers! (o|<fix-let> f val?)
  (collect-registers! (fix-let-body o) f #t)
  (collect-registers! (fix-let-arguments o) f #t))

(dm collect-registers! (o|<loop> f val?)
  (collect-registers! (loop-body o) f #t))

(export
  <box-creation>
  box-form
  <box-read>
  box-reference
  <box-write>
  <closure-creation>
  closure-creation-free
  closure-creation-index
  closurize-main!
  extract-things!
  form-definitions
  form-program
  form-quotations
  <free-reference>
  insert-box!
  gather-temporaries!
  register-allocate!
  lift!
  program-definitions
  program-form
  program-quotations
  program-bindings
  reference-offset
  reference-self?
  <renamed-local-binding>
  <top-level-form>
  flatten-seqs
  *register-passive?*
)
