;---------------------------------------------------------------------------------------------
;
;  Stack-Based Model compiler/vm, derived from
;
;  Three Implementation Models for Scheme
;  TR87-0ll
;  1987
;  R. Kent Dybvig 
;
;  https://www.cs.unc.edu/techreports/87-011.pdf
;
;
;---------------------------------------------------------------------------------------------

(load "k.sf")


;---------------------------------------------------------------------------------------------
; Utils
;---------------------------------------------------------------------------------------------

(define set-member?
  (lambda (x s)
    (cond
      [(null? s) #f]
      [(eq? x (car s)) #t]
      [else (set-member? x (cdr s))])))

(define set-cons
  (lambda (x s) 
    (if (set-member? x s) 
        s 
        (cons x s))))

(define set-union
  (lambda (sl s2)
    (if (null? sl) 
        s2 
        (set-union (cdr sl) (set-cons (car sl) s2)))))

(define set-minus
  (lambda (sl s2)
    (if (null? sl)
        '()
        (if (set-member? (car sl) s2)
            (set-minus (cdr sl) s2)
            (cons (car sl) (set-minus (cdr sl) s2))))))

(define set-intersect
  (lambda (sl s2)
    (if (null? sl)
        '()
        (if (set-member? (car sl) s2)
            (cons (car sl) (set-intersect (cdr sl) s2))
            (set-intersect (cdr sl) s2)))))

(define-syntax record-case
  (syntax-rules (else)
    [(record-case (pa . ir) clause ...)
     (let ([id (pa . ir)]) 
       (record-case id clause ...))]
    [(record-case id)
     'record-case-miss]
    [(record-case id [else exp ...])
     (begin exp ...)]
    [(record-case id [key ids exp ...] clause ...)
     (if (eq? (car id) 'key)
         (apply (lambda ids exp ...) (cdr id))
         (record-case id clause ...))]))

(define syntax-match?
  (lambda (pat exp)
    (or (eq? pat '*)
        (equal? exp pat)
        (and (pair? pat)
             (cond
              [(and (eq? (car pat) '$)
                    (pair? (cdr pat))
                    (null? (cddr pat)))
               (eq? exp (cadr pat))]
              [(and (pair? (cdr pat))
                    (eq? (cadr pat) '...)
                    (null? (cddr pat)))
               (let ([pat (car pat)])
                 (define (f lst)
                   (or (null? lst)
                       (and (pair? lst)
                            (syntax-match? pat (car lst))
                            (f (cdr lst)))))  
                 (f exp))]
              [else
               (and (pair? exp)
                    (syntax-match? (car pat) (car exp))
                    (syntax-match? (cdr pat) (cdr exp)))])))))

; unique symbol generator (poor man's version)
(define gensym
  (let ([gsc 0])
    (lambda args ; (), (symbol), or (#f) for gsc reset
      (set! gsc (fx+ gsc 1))
      (if (null? args)
          (string->symbol 
            (string-append "#" (fixnum->string gsc 10)))
          (if (symbol? (car args))
              (string->symbol 
                (string-append (symbol->string (car args))
                  (string-append "#" (fixnum->string gsc 10))))
              (set! gsc 0))))))

(define posq
  (lambda (x l)
    (let loop ([l l] [n 0])
      (cond [(null? l) #f]
            [(eq? x (car l)) n]
            [else (loop (cdr l) (fx+ n 1))]))))

(define list-diff
  (lambda (l t)
    (if (or (null? l) (eq? l t))
        '()
        (cons (car l) (list-diff (cdr l) t)))))

(define-inline (string-cmp x y)
  (%prim? "fixnum(strcmp(stringchars(obj_from_$arg), stringchars(obj_from_$arg)))" x y))

(define (pair* x . more)
  (let loop ([x x] [rest more])
    (if (null? rest) x 
        (cons x (loop (car rest) (cdr rest))))))


;---------------------------------------------------------------------------------------------
; Syntax of the Scheme Core language
;---------------------------------------------------------------------------------------------

;  <core> -> (quote <object>)
;  <core> -> (ref <id>)
;  <core> -> (set! <id> <core>)
;  <core> -> (lambda <ids> <core>) where <ids> -> (<id> ...) | (<id> ... . <id>) | <id>
;  <core> -> (begin <core> ...)
;  <core> -> (if <core> <core> <core>)
;  <core> -> (call <core> <core> ...) 

;  NB: (begin) is legit, returns unspecified value
;  on top level, these two extra core forms are legal:

;  <core> -> (define <id> <core>)
;  <core> -> (define-syntax <id> <transformer>)

; convention for 'flattened' <ids> is to put rest arg if any at the front
(define flatten-idslist
  (lambda (ilist)
    (if (list? ilist) ilist
        (let loop ([l ilist] [r '()])
          (cond [(pair? l) (loop (cdr l) (cons (car l) r))]
                [else (if (null? l) (reverse! r) (cons l (reverse! r)))])))))

(define idslist-req-count
  (lambda (ilist)
    (if (pair? ilist)
        (fx+ 1 (idslist-req-count (cdr ilist)))
        0)))


;---------------------------------------------------------------------------------------------
; Macro transformer (from Scheme to Scheme Core) derived from Al Petrofsky's EIOD 1.17
;---------------------------------------------------------------------------------------------

; An environment is a procedure that accepts any identifier and
; returns a denotation.  The denotation of an unbound identifier is
; its name (as a symbol).  A bound identifier's denotation is its
; binding, which is a pair of the current value and the identifier's 
; name (needed by quote). Biding's value can be changed later.

; Special forms are either a symbol naming a builtin, or a transformer procedure 
; that takes two arguments: a macro use and the environment of the macro use.

; <identifier>  ->  <symbol> | <thunk returning den>
; <denotation>  ->  <symbol> | <binding>
; <binding>     ->  (<symbol> . <value>)
; <value>       ->  <special> | <core>
; <special>     ->  <builtin> | <transformer>
; <builtin>     ->  syntax | define | define-syntax |
;                   quote | set! | begin | if | lambda | body
; <transformer> ->  <procedure of exp and env returning exp>

(define-inline (val-core? val)            (pair? val))
(define-inline (val-special? val)         (not (pair? val)))

(define-inline (binding? x)               (pair? x))
(define-inline (make-binding s v)         (cons s v))
(define-inline (binding-val bnd)          (cdr bnd))
(define-inline (binding-special? bnd)     (val-special? (cdr bnd)))
(define-inline (binding-sym bnd)          (car bnd))
(define-inline (binding-set-val! bnd val) (set-cdr! bnd val))
(define-inline (find-top-binding s blist) (assq s blist))

(define (new-id den)               (define p (list den)) (lambda () p))
(define (old-den id)               (car (id)))
(define (id? x)                    (or (symbol? x) (procedure? x)))
(define (id->sym id)               (if (symbol? id) id (den->sym (old-den id))))
(define (den->sym den)             (if (symbol? den) den (binding-sym den)))

(define (empty-xenv id)            (if (symbol? id) id (old-den id)))
(define (extend-xenv env id bnd)   (lambda (i) (if (eq? id i) bnd (env i))))

(define (add-binding key val env) ; adds as-is     
  (extend-xenv env key (make-binding (id->sym key) val)))

(define (add-var var val env) ; adds renamed var as <core>
  (extend-xenv env var (make-binding (id->sym var) (list 'ref val))))

; xform receives Scheme s-expressions and returns either Core Scheme <core>
; (always a pair) or special-form, which is either a builtin (a symbol) or
; a transformer (a procedure)

(define (xform appos? sexp env)
  (cond [(id? sexp) 
         (let ([hval (xform-ref sexp env)])
           (if (and (procedure? hval) (not appos?))
               (xform appos? (hval sexp env) env) ; id-syntax
               hval))]
        [(not (pair? sexp)) (xform-quote sexp env)]
        [else (let* ([head (car sexp)] [tail (cdr sexp)] [hval (xform #t head env)])
                (case hval
                  [(syntax)        (car tail)]
                  [(quote)         (xform-quote (car tail) env)]
                  [(set!)          (xform-set! (car tail) (cadr tail) env)]
                  [(begin)         (xform-begin tail env)]
                  [(if)            (xform-if tail env)]
                  [(lambda)        (xform-lambda tail env)]
                  [(body)          (xform-body tail env)]
                  [(define)        (xform-define (car tail) (cadr tail) env)]
                  [(define-syntax) (xform-define-syntax (car tail) (cadr tail) env)]
                  [else            (if (procedure? hval)
                                       (xform appos? (hval sexp env) env)
                                       (xform-call hval tail env))]))]))

(define (xform-quote sexp env)
  (list 'quote 
    (let conv ([sexp sexp])
      (cond [(id? sexp) (id->sym sexp)]
            [(pair? sexp) (cons (conv (car sexp)) (conv (cdr sexp)))]
            [(vector? sexp) (list->vector (map conv (vector->list sexp)))]
            [else sexp]))))

(define (xform-ref id env)
  (let ([den (env id)])
    (cond [(symbol? den) (list 'ref den)]
          [else (binding-val den)])))

(define (xform-set! id exp env)
  (let ([den (env id)] [xexp (xform #f exp env)])
    (cond [(symbol? den) (list 'set! den xexp)]
          [(binding-special? den) (binding-set-val! den xexp) '(begin)]
          [else (let ([val (binding-val den)])
                  (if (eq? (car val) 'ref)
                      (list 'set! (cadr val) xexp)
                      (error 'transform "set! to non-identifier form")))])))

(define (xform-begin tail env)
  (if (list? tail)
      (let ([xexps (map (lambda (sexp) (xform #f sexp env)) tail)])
        (if (and (pair? xexps) (null? (cdr xexps)))
            (car xexps) ; (begin x) => x
            (cons 'begin xexps)))
      (error 'transform "improper begin form")))

(define (xform-if tail env)
  (if (list? tail)
      (let ([xexps (map (lambda (sexp) (xform #f sexp env)) tail)])
        (case (length xexps)
          [(2) (cons 'if (append xexps '((begin))))]
          [(3) (cons 'if xexps)]
          [else (error 'transform "malformed if form")]))
      (error 'transform "improper if form")))

(define (xform-call xexp tail env)
  (if (list? tail)
      (let ([xexps (map (lambda (sexp) (xform #f sexp env)) tail)])
        (if (and (null? xexps) (eq? (car xexp) 'lambda) (null? (cadr xexp)))
            (caddr xexp) ; ((let () x)) => x
            (pair* 'call xexp xexps)))
      (error 'transform "improper application")))

(define (xform-lambda tail env)
  (if (list? tail)
      (let loop ([vars (car tail)] [ienv env] [ipars '()])
        (cond [(pair? vars)
              (let* ([var (car vars)] [nvar (gensym (id->sym var))])
                (loop (cdr vars) (add-var var nvar ienv) (cons nvar ipars)))]
              [(null? vars)
              (list 'lambda (reverse ipars) (xform-body (cdr tail) ienv))]
              [else ; improper 
              (let* ([var vars] [nvar (gensym (id->sym var))] 
                      [ienv (add-var var nvar ienv)])
                (list 'lambda (append (reverse ipars) nvar)
                  (xform-body (cdr tail) ienv)))]))
      (error 'transform "improper lambda body")))

(define (xform-body tail env)
  (if (null? tail)
      (list 'begin)
      (let loop ([env env] [ids '()] [inits '()] [nids '()] [body tail])
        (if (and (pair? body) (pair? (car body)))
            (let ([first (car body)] [rest (cdr body)])
              (let* ([head (car first)] [hval (xform #t head env)])
                (case hval
                  [(begin)
                   (loop env ids inits nids (append (cdr first) rest))]
                  [(define)
                   (let* ([id (cadr first)] [init (caddr first)]
                          [nid (gensym (id->sym id))] [env (add-var id nid env)])
                     (loop env (cons id ids) (cons init inits) (cons nid nids) rest))]
                  [(define-syntax)
                   (let* ([id (cadr first)] [init (caddr first)]
                          [env (add-binding id '(undefined) env)])
                     (loop env (cons id ids) (cons init inits) (cons #t nids) rest))]
                  [else
                   (if (procedure? hval)
                       (loop env ids inits nids (cons (hval first env) rest))
                       (xform-labels (reverse ids) (reverse inits) (reverse nids) body env))])))
            (xform-labels (reverse ids) (reverse inits) (reverse nids) body env)))))

(define (xform-labels ids inits nids body env)
  (let loop ([ids ids] [inits inits] [nids nids] [sets '()] [lids '()])
    (cond [(null? ids) 
           (let* ([xexps (append (reverse sets)
                           (map (lambda (sexp) (xform #f sexp env)) body))]
                  [xexp (if (and (pair? xexps) (null? (cdr xexps)))
                            (car xexps)
                            (cons 'begin xexps))])
             (if (null? lids)
                 xexp
                 (pair* 'call (list 'lambda (reverse lids) xexp)
                   (map (lambda (lid) '(begin)) lids))))]
          [(symbol? (car nids)) ; define
           (loop (cdr ids) (cdr inits) (cdr nids)
             (cons (xform-set! (car ids) (car inits) env) sets)
             (cons (car nids) lids))]
          [else ; define-syntax
           (binding-set-val! (env (car ids)) (xform #t (car inits) env))
           (loop (cdr ids) (cdr inits) (cdr nids) sets lids)])))

(define (xform-define id exp env) ; top-level only
  (if (id? id)
      (list 'define (id->sym id) (xform #f exp env))
      (error 'transform "define of non-identifier form")))

(define (xform-define-syntax id exp env) ; top-level only
  (if (id? id)
      (list 'define-syntax (id->sym id) (xform #t exp env))
      (error 'transform "define-syntax of non-identifier form")))


; ellipsis denotation is used for comparisons only

(define denotation-of-default-ellipsis 
  (make-binding '... (lambda (sexp env) (error '... sexp))))

(define *transformers* 
  (list
    (make-binding 'syntax 'syntax) 
    (make-binding 'define 'define)
    (make-binding 'define-syntax 'define-syntax)
    (make-binding 'quote 'quote)
    (make-binding 'set! 'set!)
    (make-binding 'begin 'begin)
    (make-binding 'if 'if)
    (make-binding 'lambda 'lambda)
    (make-binding 'body 'body)
    denotation-of-default-ellipsis))

(define (top-transformer-env id)
  (let ([bnd (find-top-binding id *transformers*)])
    (cond [(binding? bnd)
           ; special case: syntax-rules in sexp form (left by init)
           (let ([val (binding-val bnd)])
             (when (and (pair? val) (eq? (car val) 'syntax-rules))
               (binding-set-val! bnd (transform #t val))))
           bnd]
          [(symbol? id)
           (let ([bnd (make-binding id (list 'ref id))])
             (set! *transformers* (cons bnd *transformers*))
             bnd)]
          [else (old-den id)])))

(define (install-transformer! s t)
  (binding-set-val! (top-transformer-env s) t))

(define (install-transformer-rules! s ell lits rules)
  (install-transformer! s
    (syntax-rules* top-transformer-env ell lits rules)))

(define (transform appos? sexp . optenv)
  (gensym #f) ; reset gs counter to make results reproducible
  (xform appos? sexp (if (null? optenv) top-transformer-env (car optenv))))


; 'syntax-rules' transformer produces another transformer from the rules

(define (syntax-rules* mac-env ellipsis pat-literals rules)
  (define (pat-literal? id) (memq id pat-literals))
  (define (not-pat-literal? id) (not (pat-literal? id)))
  (define (ellipsis-pair? x)
    (and (pair? x) (ellipsis? (car x))))
  (define (ellipsis? x)
    (if ellipsis
        (eq? x ellipsis)
        (and (id? x) (eq? (mac-env x) denotation-of-default-ellipsis))))

  ; List-ids returns a list of the non-ellipsis ids in a
  ; pattern or template for which (pred? id) is true.  If
  ; include-scalars is false, we only include ids that are
  ; within the scope of at least one ellipsis.
  (define (list-ids x include-scalars pred?)
    (let collect ([x x] [inc include-scalars] [l '()])
      (cond [(id? x) (if (and inc (pred? x)) (cons x l) l)]
            [(vector? x) (collect (vector->list x) inc l)]
            [(pair? x)
             (if (ellipsis-pair? (cdr x))
                 (collect (car x) #t (collect (cddr x) inc l))
                 (collect (car x) inc (collect (cdr x) inc l)))]
            [else l])))

  ; Returns #f or an alist mapping each pattern var to a part of
  ; the input.  Ellipsis vars are mapped to lists of parts (or
  ; lists of lists ...).
  (define (match-pattern pat use use-env)
    (call-with-current-continuation
      (lambda (return)
        (define (fail) (return #f))
        (let match ([pat pat] [sexp use] [bindings '()])
          (define (continue-if condition)
            (if condition bindings (fail)))
          (cond
            [(id? pat)
              (if (pat-literal? pat)
                  (continue-if (and (id? sexp) (eq? (use-env sexp) (mac-env pat))))
                  (cons (cons pat sexp) bindings))]
            [(vector? pat)
             (or (vector? sexp) (fail))
             (match (vector->list pat) (vector->list sexp) bindings)]
            [(not (pair? pat))
             (continue-if (equal? pat sexp))]
            [(ellipsis-pair? (cdr pat))
             (let* ([tail-len (length (cddr pat))]
                    [sexp-len (if (list? sexp) (length sexp) (fail))]
                    [seq-len (fx- sexp-len tail-len)]
                    [sexp-tail (begin (if (negative? seq-len) (fail)) (list-tail sexp seq-len))]
                    [seq (reverse (list-tail (reverse sexp) tail-len))]
                    [vars (list-ids (car pat) #t not-pat-literal?)])
                     (define (match1 sexp)
                       (map cdr (match (car pat) sexp '())))
                     (append
                       (apply map (cons list (cons vars (map match1 seq))))
                       (match (cddr pat) sexp-tail bindings)))]
            [(pair? sexp)
             (match (car pat) (car sexp)
               (match (cdr pat) (cdr sexp) bindings))]
            [else (fail)])))))

  (define (expand-template pat tmpl top-bindings)
    ; New-literals is an alist mapping each literal id in the
    ; template to a fresh id for inserting into the output.  It
    ; might have duplicate entries mapping an id to two different
    ; fresh ids, but that's okay because when we go to retrieve a
    ; fresh id, assq will always retrieve the first one.
    (define new-literals
      (map (lambda (id) (cons id (new-id (mac-env id))))
           (list-ids tmpl #t
             (lambda (id) (not (assq id top-bindings))))))

    (define ellipsis-vars
      (list-ids pat #f not-pat-literal?))

    (define (list-ellipsis-vars subtmpl)
      (list-ids subtmpl #t 
        (lambda (id) (memq id ellipsis-vars))))

    (let expand ([tmpl tmpl] [bindings top-bindings])
      (let expand-part ([tmpl tmpl])
        (cond
          [(id? tmpl)
           (cdr (or (assq tmpl bindings)
                    (assq tmpl top-bindings)
                    (assq tmpl new-literals)))]
          [(vector? tmpl)
           (list->vector (expand-part (vector->list tmpl)))]
          [(pair? tmpl)
           (if (ellipsis-pair? (cdr tmpl))
               (let ([vars-to-iterate (list-ellipsis-vars (car tmpl))])
                 (define (lookup var)
                   (cdr (assq var bindings)))
                 (define (expand-using-vals . vals)
                   (expand (car tmpl)
                     (map cons vars-to-iterate vals)))
                 (let ([val-lists (map lookup vars-to-iterate)])
                   (append
                     (apply map (cons expand-using-vals val-lists))
                     (expand-part (cddr tmpl)))))
               (cons (expand-part (car tmpl)) (expand-part (cdr tmpl))))]
          [else tmpl]))))

  (lambda (use use-env)
    (let loop ([rules rules])
      (if (null? rules) (error 'transform "invalid syntax" use))
      (let* ([rule (car rules)] [pat (car rule)] [tmpl (cadr rule)])
        (cond [(match-pattern pat use use-env) =>
               (lambda (bindings) (expand-template pat tmpl bindings))]
              [else (loop (cdr rules))])))))

(install-transformer! 'syntax-rules
  (lambda (sexp env)
    (define syntax-id (new-id (make-binding 'syntax 'syntax)))
    ; sexp can be either
    (if (id? (cadr sexp))
        ; (_ ellipsis (litname ...) . rules)
        (list syntax-id (syntax-rules* env (cadr sexp) (caddr sexp) (cdddr sexp)))
        ; or (_ (litname ...) . rules)
        (list syntax-id (syntax-rules* env #f (cadr sexp) (cddr sexp))))))

; non-recursive transformer for define relies on old definition

(install-transformer! 'define
  (let ([env (add-binding 'define 'define top-transformer-env)])
    (syntax-rules* env #f '() '(
      [(_ (name . args) . forms)
       (define name (lambda args . forms))]
      [(_ name exp)
       (define name exp)]))))

; Remaining transformers are made with the help of syntax-rules*
; NB: order of installation is important -- each transformer can
; be self-recursive but can't use transformers defined later!

(define-syntax install-sr-transformer!
  (syntax-rules (quote syntax-rules)
    [(_ 'name (syntax-rules (lit ...) . rules))
     (install-transformer-rules! 'name #f '(lit ...) 'rules)]
    [(_ 'name (syntax-rules ellipsis (lit ...) . rules))
     (install-transformer-rules! 'name 'ellipsis '(lit ...) 'rules)]))    

(install-sr-transformer! 'letrec-syntax
  (syntax-rules ()
    [(_ ([key trans] ...) . forms) ; non-splicing!
     (body (define-syntax key trans) ... . forms)]))

(install-sr-transformer! 'let-syntax
  (syntax-rules ()
    [(_ () . forms) 
     (body . forms)]
    [(_ ([key trans] . bindings) . forms)
     (letrec-syntax ([temp trans])
       (let-syntax bindings 
         (letrec-syntax ([key temp]) . forms)))]))

(install-sr-transformer! 'letrec
  (syntax-rules ()
    [(_ ([var init] ...) . forms)
     (body (define var init) ... . forms)]))

(install-sr-transformer! 'let
  (syntax-rules ()
    [(_ ([var init] ...) . forms)
     ((lambda (var ...) . forms) init ...)]
    [(_ name ([var init] ...) . forms)
     ((letrec ((name (lambda (var ...) . forms))) name) init ...)]))

(install-sr-transformer! 'let*
  (syntax-rules ()
    [(_ () . forms) 
     (body . forms)]
    [(_ (first . more) . forms)
     (let (first) (let* more . forms))]))

(install-sr-transformer! 'and
  (syntax-rules ()
    [(_) #t]
    [(_ test) test]
    [(_ test . tests) (if test (and . tests) #f)]))

(install-sr-transformer! 'or
  (syntax-rules ()
    [(_) #f]
    [(_ test) test]
    [(_ test . tests) (let ([x test]) (if x x (or . tests)))]))

(install-sr-transformer! 'cond
  (syntax-rules (else =>)
    [(_) #f]
    [(_ (else . exps)) (begin . exps)]
    [(_ (x) . rest) (or x (cond . rest))]
    [(_ (x => proc) . rest) (let ([tmp x]) (cond [tmp (proc tmp)] . rest))]
    [(_ (x . exps) . rest) (if x (begin . exps) (cond . rest))]))

(install-sr-transformer! 'case-test
  (syntax-rules (else) 
    [(_ k else) #t]
    [(_ k atoms) (memv k 'atoms)]))

(install-sr-transformer! 'case
  (syntax-rules ()
    [(_ x (test . exprs) ...)
     (let ([key x]) (cond ((case-test key test) . exprs) ...))]))

(install-sr-transformer! 'do
  (syntax-rules ()
    [(_ ((var init . step) ...) ending expr ...)
     (let loop ([var init] ...)
       (cond ending [else expr ... (loop (begin var . step) ...)]))]))

(install-sr-transformer! 'quasiquote
  (syntax-rules (unquote unquote-splicing quasiquote)
    [(_ ,x) x]
    [(_ (,@x . y)) (append x `y)]
    [(_ `x . d) (cons 'quasiquote (quasiquote (x) d))]
    [(_ ,x   d) (cons 'unquote (quasiquote (x) . d))]
    [(_ ,@x  d) (cons 'unquote-splicing (quasiquote (x) . d))]
    [(_ (x . y) . d) (cons (quasiquote x . d) (quasiquote y . d))]
    [(_ #(x ...) . d) (list->vector (quasiquote (x ...) . d))]
    [(_ x . d) 'x]))

(install-sr-transformer! 'delay
  (syntax-rules ()
    [(_ exp)
     (make-delayed (lambda () exp))]))

(install-sr-transformer! 'when
  (syntax-rules ()
    [(_ test . rest) (if test (begin . rest))]))

(install-sr-transformer! 'unless
  (syntax-rules ()
    [(_ test . rest) (if (not test) (begin . rest))]))

;(install-sr-transformer! 'define-inline
;  (syntax-rules ()
;    [(_ . rest) (define . rest)]))


;---------------------------------------------------------------------------------------------
; Runtime
;---------------------------------------------------------------------------------------------

(%localdef "#include \"i.h\"")

(define *globals* '())

(define global-location
  (lambda (sym)
    (let ([loc (assq sym *globals*)])
      (if (pair? loc)
          loc
          (let ([loc (cons sym 'undefined)])
            (set! *globals* (cons loc *globals*))
            loc)))))

(define-syntax index-global cdr)
(define-syntax index-set-global! set-cdr!)


;---------------------------------------------------------------------------------------------
; String representation of S-expressions and code arguments
;---------------------------------------------------------------------------------------------

(define (write-serialized-char x port)
  (cond [(or (char=? x #\%) (char=? x #\") (char=? x #\\) (char<? x #\space) (char>? x #\~))
         (write-char #\% port)
         (let ([s (fixnum->string (char->integer x) 16)])
           (if (fx=? (string-length s) 1) (write-char #\0 port))
           (write-string s port))]
        [else (write-char x port)]))

(define (write-serialized-size n port)         
  (write-string (fixnum->string n 10) port)
  (write-char #\: port))

(define (write-serialized-element x port)
  (write-serialized-sexp x port)
  (write-char #\; port))

(define (write-serialized-sexp x port)
  (cond [(eq? x #f) 
         (write-char #\f port)]
        [(eq? x #t) 
         (write-char #\t port)]
        [(eq? x '()) 
         (write-char #\n port)]
        [(char? x)
         (write-char #\c port)
         (write-serialized-char x port)]
        [(number? x)
         (write-char (if (exact? x) #\i #\j) port)
         (write-string (number->string x 10) port)]
        [(list? x)
         (write-char #\l port)
         (write-serialized-size (length x) port)
         (do ([x x (cdr x)]) [(null? x)]
           (write-serialized-element (car x) port))]
        [(pair? x)
         (write-char #\p port)
         (write-serialized-element (car x) port) 
         (write-serialized-element (cdr x) port)]
        [(vector? x)
         (write-char #\v port)
         (write-serialized-size (vector-length x) port)
         (do ([i 0 (fx+ i 1)]) [(fx=? i (vector-length x))]
           (write-serialized-element (vector-ref x i) port))]
        [(string? x)
         (write-char #\s port)
         (write-serialized-size (string-length x) port)
         (do ([i 0 (fx+ i 1)]) [(fx=? i (string-length x))]
           (write-serialized-char (string-ref x i) port))]
        [(symbol? x)
         (write-char #\y port)
         (let ([x (symbol->string x)])
           (write-serialized-size (string-length x) port)
           (do ([i 0 (fx+ i 1)]) [(fx=? i (string-length x))]
             (write-serialized-char (string-ref x i) port)))]
        [else (error 'encode-sexp "cannot encode literal: ~s" x)]))

(define (write-serialized-arg arg port)
  (if (and (number? arg) (exact? arg) (fx<=? 0 arg) (fx<=? arg 9))
      (write-char (string-ref "0123456789" arg) port)
      (begin (write-char #\( port)
            (write-serialized-sexp arg port)
            (write-char #\) port))))


;---------------------------------------------------------------------------------------------
; Compiler producing serialized code
;---------------------------------------------------------------------------------------------

(define find-free*
  (lambda (x* b)
    (if (null? x*)
        '()
        (set-union 
          (find-free (car x*) b) 
          (find-free* (cdr x*) b)))))

(define find-free
  (lambda (x b)
    (record-case x
      [quote (obj) 
       '()]
      [ref (id)
       (if (set-member? id b) '() (list id))]
      [set! (id exp)
       (set-union
         (if (set-member? id b) '() (list id))
         (find-free exp b))]
      [lambda (idsi exp)
       (find-free exp (set-union (flatten-idslist idsi) b))]
      [if (test then else)
       (set-union
         (find-free test b)
         (set-union (find-free then b) (find-free else b)))]
      [begin exps
       (find-free* exps b)]
      [call (exp . args)
       (set-union (find-free exp b) (find-free* args b))])))

(define find-sets*
  (lambda (x* v)
    (if (null? x*)
        '()
        (set-union 
          (find-sets (car x*) v) 
          (find-sets* (cdr x*) v)))))

(define find-sets
  (lambda (x v)
    (record-case x
      [quote (obj) 
       '()]
      [ref (id)
       '()]
      [set! (id x)
       (set-union
         (if (set-member? id v) (list id) '())
         (find-sets x v))]
      [lambda (idsi exp)
       (find-sets exp (set-minus v (flatten-idslist idsi)))]
      [begin exps
       (find-sets* exps v)]
      [if (test then else)
       (set-union
         (find-sets test v)
         (set-union (find-sets then v) (find-sets else v)))]
      [call (exp . args)
       (set-union (find-sets exp v) (find-sets* args v))])))


(define find-integrable-encoding
  (%prim "{ /* define find-integrable-encoding */
    static obj c[] = { obj_from_objptr(vmcases+4) };
    $return objptr(c); }"))

(define encode-integrable
  (%prim "{ /* define encode-integrable */
    static obj c[] = { obj_from_objptr(vmcases+5) };
    $return objptr(c); }"))


(define codegen
  ; x: Scheme Core expression to compile
  ; l: local var list (with #f placeholders for nonvar slots)
  ; f: free var list
  ; s: set! var set
  ; g: global var set
  ; k: #f: x goes to ac, N: x is to be returned after (sdrop n)
  ; port: output code goes here
  (lambda (x l f s g k port)
    (record-case x
      [quote (obj)
       (case obj
          [(#t) (write-char #\t port)]
          [(#f) (write-char #\f port)]
          [(()) (write-char #\n port)]
          [else (write-char #\' port) (write-serialized-arg obj port)])
       (when k (write-char #\] port) (write-serialized-arg k port))]
      [ref (id)
       (cond [(posq id l) => ; local
              (lambda (n) 
                (write-char #\. port)
                (write-serialized-arg n port)
                (if (set-member? id s) (write-char #\^ port)))]
             [(posq id f) => ; free
              (lambda (n) 
                (write-char #\: port)
                (write-serialized-arg n port)
                (if (set-member? id s) (write-char #\^ port)))]
             [else ; global
              (write-char #\@ port)
              (write-serialized-arg id port)])
       (when k (write-char #\] port) (write-serialized-arg k port))]
      [set! (id x)
       (codegen x l f s g #f port)
       (cond [(posq id l) => ; local
              (lambda (n) 
                (write-char #\. port) (write-char #\! port)
                (write-serialized-arg n port))]
             [(posq id f) => ; free
              (lambda (n) 
                (write-char #\: port) (write-char #\! port)
                (write-serialized-arg n port))]
             [else ; global
              (write-char #\@ port) (write-char #\! port)
              (write-serialized-arg id port)])
       (when k (write-char #\] port) (write-serialized-arg k port))]
      [begin exps
       (let loop ([xl exps])
         (when (pair? xl)
           (let ([k (if (pair? (cdr xl)) #f k)])
             (codegen (car xl) l f s g k port)
             (loop (cdr xl)))))
       (when (and k (null? exps)) (write-char #\] port) (write-serialized-arg k port))]
      [if (test then else)
       (codegen test l f s g #f port)
       (write-char #\? port)
       (write-char #\{ port)
       (codegen then l f s g k port)
       (write-char #\} port)
       (cond [k ; tail call: 'then' arm exits, so br around is not needed
              (codegen else l f s g k port)]
             [(equal? else '(begin)) ; non-tail with void 'else' arm
              ] ; no code needed -- ac retains #f from failed test
             [else ; non-tail with 'else' expression; needs br 
              (write-char #\{ port)
              (codegen else l f s g k port)
              (write-char #\} port)])]              
      [lambda (idsi exp)
       (let* ([ids (flatten-idslist idsi)]
              [free (set-minus (find-free exp ids) g)]
              [sets (find-sets exp ids)])
         (do ([free (reverse free) (cdr free)] [l l (cons #f l)]) [(null? free)]
           ; note: called with empty set! var list
           ; to make sure no dereferences are generated
           (codegen (list 'ref (car free)) l f '() g #f port)
           (write-char #\, port))
         (write-char #\& port)
         (write-serialized-arg (length free) port)
         (write-char #\{ port)
         (cond [(list? idsi)
                (write-char #\% port)
                (write-serialized-arg (length idsi) port)]
               [else
                (write-char #\% port) (write-char #\! port)
                (write-serialized-arg (idslist-req-count idsi) port)])
         (do ([ids ids (cdr ids)] [n 0 (fx+ n 1)]) [(null? ids)]
           (when (set-member? (car ids) sets)
             (write-char #\# port)
             (write-serialized-arg n port)))
         (codegen exp ids free 
           (set-union sets (set-intersect s free))
           g (length ids) port)
         (write-char #\} port))
       (when k (write-char #\] port) (write-serialized-arg k port))]
      [call (exp . args)
       (cond [(and (eq? (car exp) 'lambda) (list? (cadr exp))
                   (fx=? (length args) (length (cadr exp))))
              ; let-like call; compile as special lambda + call combo
              (do ([args (reverse args) (cdr args)] [l l (cons #f l)]) 
                [(null? args)]
                (codegen (car args) l f s g #f port)
                (write-char #\, port))
              (let* ([ids (cadr exp)] [exp (caddr exp)]
                     [sets (find-sets exp ids)]
                     [news (set-union (set-minus s ids) sets)]
                     [newl (append ids l)]) ; with real names
                (do ([ids ids (cdr ids)] [n 0 (fx+ n 1)]) [(null? ids)]
                  (when (set-member? (car ids) sets)
                    (write-char #\# port)
                    (write-serialized-arg n port)))
                (if k 
                    (codegen exp newl f news g (fx+ k (length args)) port)
                    (begin 
                      (codegen exp newl f news g #f port)
                      (write-char #\_ port)
                      (write-serialized-arg (length args) port))))]
             [(and (eq? (car exp) 'ref) 
                   (not (posq (cadr exp) l)) (not (posq (cadr exp) f)) 
                   (find-integrable-encoding (cadr exp) (length args))) =>
              ; integrable function/procedure
              (lambda (ienc)
                ; regular convention is 1st arg in a, others on stack
                (do ([args (reverse args) (cdr args)] [l l (cons #f l)]) 
                  [(null? args)]
                  (codegen (car args) l f s g #f port)
                  (unless (null? (cdr args)) (write-char #\, port)))
                (encode-integrable (length args) ienc port)
                (when k (write-char #\] port) (write-serialized-arg k port)))]
             [k 
              ; tail call with k elements under arguments
              (do ([args (reverse args) (cdr args)] [l l (cons #f l)]) 
                [(null? args) (codegen exp l f s g #f port)]
                (codegen (car args) l f s g #f port)
                (write-char #\, port))
              (write-char #\[ port)
              (write-serialized-arg k port)
              (write-serialized-arg (length args) port)]
             [else
              ; non-tail call; 'save' puts 2 extra elements on the stack!
              (write-char #\$ port) (write-char #\{ port)
              (do ([args (reverse args) (cdr args)] [l (cons #f (cons #f l)) (cons #f l)]) 
                [(null? args) (codegen exp l f s g #f port)]
                (codegen (car args) l f s g #f port)
                (write-char #\, port))
              (write-char #\[ port)
              (write-serialized-arg 0 port)
              (write-serialized-arg (length args) port)
              (write-char #\} port)])])))

(define (compile-to-string x)
  (let ([p (open-output-string)])
    (codegen x '() '() '() (find-free x '()) #f p)                   
    (get-output-string p)))


;---------------------------------------------------------------------------------------------
; Code deserializer and Evaluator (use built-ins)
;---------------------------------------------------------------------------------------------

(define execute-thunk-closure
  (%prim "{ /* define execute-thunk-closure */
    static obj c[] = { obj_from_objptr(vmcases+0) }; 
    $return objptr(c); }"))

(define make-closure
  (%prim "{ /* define make-closure */
    static obj c[] = { obj_from_objptr(vmcases+1) };
    $return objptr(c); }"))

(define execute
  (lambda (code)
    (execute-thunk-closure (make-closure code))))

(define decode-sexp
  (%prim "{ /* define decode-sexp */
    static obj c[] = { obj_from_objptr(vmcases+2) };
    $return objptr(c); }"))

(define decode
  (%prim "{ /* define decode */
    static obj c[] = { obj_from_objptr(vmcases+3) };
    $return objptr(c); }"))

(define (evaluate x)
  (execute (decode (compile-to-string (transform #f x)))))


;---------------------------------------------------------------------------------------------
; File processor (Scheme => Serialized code)
;---------------------------------------------------------------------------------------------

;(define *defined-refs* '())
;(define *free-refs* '())
(define *hide-refs* '(
  define-inline nullary-unary-adaptor nullary-unary-binary-adaptor
  unary-binary-adaptor unary-binary-ternary-adaptor 
  unary-binary-ternary-quaternary-adaptor binary-ternary-adaptor
  cmp-reducer addmul-reducer subdiv-reducer append-reducer
))

(define (display-code cstr oport)
  (let loop ([i 0] [l (string-length cstr)])
    (let ([r (fx- l i)])
      (cond [(<= r 70)
             (display "  \"" oport)
             (display (substring cstr i l))
             (display "\"," oport)]
            [else
             (display "  \"" oport)
             (display (substring cstr i (fx+ i 70)))
             (display "\"\n" oport)             
             (loop (fx+ i 70) l)]))))  


(define (process-define-syntax id xval oport)
  (newline oport)
  (display "  \"" oport) (display id oport) (display "\",\n" oport)
  ; hack xval's define-inline leftovers
  (set! xval 
    (let hack ([v xval])
      (cond [(procedure? v) 'syntax-rules]
            [(eq? v 'define-inline) '_]
            [(pair? v) (cons (hack (car v)) (hack (cdr v)))]
            [else v])))
  (let ([p (open-output-string)]) (write-serialized-sexp xval p)
    (display-code (get-output-string p) oport) (newline oport)))

(define (process-statement xval oport)
  (define cstr (compile-to-string xval))
  (newline oport)
  (display "  0,\n" oport)
  (display-code cstr oport) (newline oport))

(define (process-define id xlam oport)
  ;(newline oport)
  ;(display "  \"P\", \"" oport) (display id oport) (display "\",\n" oport)
  ;(display-code (compile-to-string xlam) oport) (newline oport)
  (process-statement (list 'set! id xlam) oport))

(define (process-top-form x oport)
  (cond 
    [(pair? x)
     (let ([hval (transform #t (car x))])
       (cond
         [(eq? hval 'begin)
          (let loop ([x* (cdr x)])
            (when (pair? x*) 
               (process-top-form (car x*) oport)
               (loop (cdr x*))))]
         [(eq? hval 'define-syntax)
          (let ([xval (transform #t (caddr x))])
            (install-transformer! (cadr x) xval)
            (unless (memq (cadr x) *hide-refs*)  
              (process-define-syntax (cadr x) (caddr x) oport)))]
         [(eq? hval 'define)
          (let ([xval (transform #f (caddr x))])
            (process-define (cadr x) xval oport))]
         [(procedure? hval)
          (process-top-form (hval x top-transformer-env) oport)]
         [else
          (process-statement (transform #f x) oport)]))]
    [else
     (process-statement (transform #f x) oport)]))

(define (path-strip-directory filename)
  (let loop ([l (reverse (string->list filename))] [r '()])
    (cond [(null? l) (list->string r)]
          [(memv (car l) '(#\\ #\/ #\:)) (list->string r)]
          [else (loop (cdr l) (cons (car l) r))])))

(define (path-strip-extension filename)
  (let ([l (reverse (string->list filename))])
    (let ([r (memv #\. l)])
      (if r (list->string (reverse (cdr r))) filename))))

(define (module-name filename)
  (path-strip-extension (path-strip-directory filename)))

(define (process-file fname)
  (define iport (open-input-file fname))
  (define oport (current-output-port))
  (define mname (module-name fname))
  (display "/* " oport) (display mname oport) 
  (display ".c -- generated via skint -c " oport) 
  (display (path-strip-directory fname) oport) 
  (display " */" oport) (newline oport) (newline oport)
  (display "char *" oport) (display mname oport) 
  (display "_code[] = {" oport) (newline oport)
  (let loop ([x (read iport)])
    (unless (eof-object? x)
      (process-top-form x oport)
      (loop (read iport))))
  (display "\n  0, 0\n};\n" oport)
  (close-input-port iport))


;---------------------------------------------------------------------------------------------
; Test environment
;---------------------------------------------------------------------------------------------

; NB: 'nuate' restores stack with fn arg on top of return triple
(define continuation-closure-code (decode "%1.0K2]1"))

#|
(define (install-global-lambdas)
  (define (install-global-lambda! sym cstr)
    (index-set-global! (global-location sym) (make-closure (decode cstr))))
  ; top-level definitions using integrables 
  (install-global-lambda! 'eq? "%2.1,.1q]2")
  (install-global-lambda! 'car "%1.0a]1")
  (install-global-lambda! 'cdr "%1.0d]1")
  (install-global-lambda! 'null? "%1.0u]1")
  (install-global-lambda! 'pair? "%1.0p]1")
  (install-global-lambda! 'cons "%2.1,.1c]2")
  (install-global-lambda! 'list "%!0.0]1")
  (install-global-lambda! 'append "%2.1,.1L6]2")
  (install-global-lambda! '+ "%2.1,.1+]2")
  (install-global-lambda! '- "%2.1,.1-]2")
  (install-global-lambda! '* "%2.1,.1*]2")
  (install-global-lambda! '< "%2.1,.1<]2")
  (install-global-lambda! '= "%2.1,.1=]2")
  (install-global-lambda! 'zero? "%1.0z]1")
  (install-global-lambda! 'not "%1.0~]1")
  (install-global-lambda! 'string<? "%2'0,.2,.2O6<]2")
  (install-global-lambda! 'deserialize-sexp "%1.0U3]1")
  (install-global-lambda! 'deserialize-code "%1.0U4]1")
  (install-global-lambda! 'call/cc "%1K1,.1[11"))
|#

(define install-global-lambdas
  (%prim "{ /* define install-global-lambdas */
    static obj c[] = { obj_from_objptr(vmcases+6) };
    $return objptr(c); }"))

(install-global-lambdas)


;---------------------------------------------------------------------------------------------
; Tests 
;---------------------------------------------------------------------------------------------

(define test1
 '(let ()
    (define (sort-list obj pred)
      (define (loop l)
        (if (and (pair? l) (pair? (cdr l))) (split l '() '()) l))
      (define (split l one two)
        (if (pair? l)
            (split (cdr l) two (cons (car l) one))
            (merge (loop one) (loop two))))
      (define (merge one two)
        (cond
          [(null? one) two]
          [(pred (car two) (car one))
           (cons (car two) (merge (cdr two) one))]
          [else (cons (car one) (merge (cdr one) two))]))
      (loop obj))
    (sort-list
      '("one" "two" "three" "four" "five" "six"
        "seven" "eight" "nine" "ten" "eleven" "twelve")
      string<?)))

(define test2
 '(let ()
    (define tak
      (lambda (x y z)
        (if (< y x)
            (tak (tak (- x 1) y z)
                (tak (- y 1) z x)
                (tak (- z 1) x y))
            z)))
    (define runtak
      (lambda (n r)
        (let loop ([n n] [r r] [s 7])
          (if (= n 0) r
              (let ([v (tak 18 12 (- s 1))])
                (loop (- n 1) (+ r v) v))))))
    (runtak 10 0)))

(define test3
 '(let ()
    (define (nqueens n)
      (define (one-to n)
        (let loop ((i n) (l '()))
          (cond
            ((zero? i) l)
            (else (loop (- i 1) (cons i l))))))
      (define (try-it x y z)
        (if (null? x)
            (if (null? y) 1 0)
            (+ (if (ok? (car x) 1 z)
                   (try-it (append (cdr x) y) '() (cons (car x) z))
                   0)
              (try-it (cdr x) (cons (car x) y) z))))
      (define (ok? row dist placed)
        (if (null? placed) #t
            (and (not (= (car placed) (+ row dist)))
                 (not (= (car placed) (- row dist)))
                 (ok? row (+ dist 1) (cdr placed)))))
      (try-it (one-to n) '() '()))
    (define (run-test count)
      (let loop ((n count) (v 92))
        (cond
          ((zero? n) v)
          (else (loop (- n 1) (nqueens (- v 84)))))))
    (run-test 10)))

(define test4
 '(let ()
    (define y
      (lambda (e) 
        ((call/cc call/cc) 
        (lambda (f) 
          (e (lambda (x) (((call/cc (call/cc call/cc)) f) x)))))))
    (define fakt
      (y (lambda (self) (lambda (x) (if (= x 0) 1 (* x (self (- x 1))))))))
    (fakt 10)))

(define test5
 '(let ()
    (define y
      (lambda (e) 
        ((call/cc call/cc) 
        (lambda (f) 
          (e (lambda (x) (((call/cc (call/cc call/cc)) f) x)))))))
    (define fakty
      (y (lambda (self) 
            (lambda (x) (if (= x 0) 1 (* x (self (- x 1))))))))
    (define (fakti x)
      (let loop ((n 1) (x x))
          (if (= x 1)
            n
            (loop (* n x) (- x 1)))))
    (define (faktr x)
      (if (= x 1)
          1
          (* x (faktr (- x 1)))))
    (define faktl
      (lambda (x)
          ((lambda (self) (self self x))
          (lambda (self x)
            (if (= x 1)
                x
                (* (self self (- x 1)) x))))))
    (let ([y (fakty 10)] [i (fakti 10)] [r (faktr 10)] [l (faktl 10)])
      (cons y (cons i (cons r (cons l '())))))))

; (evaluate test1) =>
; ("eight" "eleven" "five" "four" "nine" "one" "seven" "six" "ten" "three" "twelve" "two")     
;
; (evaluate test2) =>
; 70
;
; (evaluate test3) =>
; 92
;
; (evaluate test4) =>
; 3628800
;
; (evaluate test5) =>
; (3628800 3628800 3628800 3628800)
;


;---------------------------------------------------------------------------------------------
; REPL
;---------------------------------------------------------------------------------------------

(define (run-tests)
  (define start (current-jiffy))
  (display "Running tests ...") (newline)
  (write (evaluate test1)) (newline)
  (write (evaluate test2)) (newline)
  (write (evaluate test3)) (newline)
  (write (evaluate test4)) (newline)
  (write (evaluate test5)) (newline)
  (display "Elapsed time: ") (write (* 1000 (/ (- (current-jiffy) start) (jiffies-per-second))))
  (display " ms.") (newline))

(define (repl-eval x)
  (let ([xexp (transform #f x)])
    (display "TRANSFORM =>") (newline)
    (write xexp) (newline)
    (if (eq? (car xexp) 'define) (set-car! xexp 'set!))
    (display "COMPILE-TO-STRING =>") (newline)
    (let ([cstr (compile-to-string xexp)] [start #f])
      (display cstr) (newline)
      (display "DECODE+EXECUTE =>") (newline) 
      (set! start (current-jiffy))
      (let* ([thunk (decode cstr)]
             [foo (begin (display "decoded: ") (write thunk) (newline))]
             [res (execute thunk)])
        (write res) (newline))
      (display "Elapsed time: ") (write (* 1000 (/ (- (current-jiffy) start) (jiffies-per-second))))
      (display " ms.") (newline))))

(define (repl-eval-top-form x)
  (cond 
    [(and (list? x) (= (length x) 2) (eq? (car x) 'load) (string? (cadr x)))
     (let ([iport (open-input-file (cadr x))])
       (repl-from-port iport)
       (close-input-port iport))]
    [(pair? x)
     (let ([hval (transform #t (car x))])
       (cond
         [(eq? hval 'begin)
          (let loop ([x* (cdr x)])
            (when (pair? x*) 
              (repl-eval-top-form (car x*))
              (loop (cdr x*))))]
         [(eq? hval 'define-syntax)
          (let ([xval (transform #t (caddr x))])
            (install-transformer! (cadr x) xval))]
         [(procedure? hval)
          (repl-eval-top-form (hval x top-transformer-env))]
         [else
          (repl-eval x)]))]
    [else
     (repl-eval x)]))

(define (repl-read iport)
  (when (eq? iport (current-input-port))
    (display "\n3imp> "))
  (read iport))

(define (repl-from-port iport)
  (let loop ([x (repl-read iport)])
    (unless (eof-object? x)
        (repl-eval-top-form x)
        (loop (repl-read iport)))))

(define (run-repl)
  (repl-from-port (current-input-port)))

(define (main argv)
  (let ([args (cdr (command-line))])
    (cond
      [(syntax-match? '("-c" *) args)
       (process-file (cadr args))]
      [else
       ;(run-tests)
       (run-repl)])))