(use-modules (ice-9 match)
             (srfi srfi-1) 
             (srfi srfi-9) ;; Record types
             (srfi srfi-64) ;; Testing framework
             )
;; Compiler o'clock!

;; Globals
(define undefined-value (make-symbol "undefined"))
(define leafs '())
(define compounds '())

;; registers
(define *env* (make-parameter #f))
(define *val* (make-parameter #f))
(define *fun* (make-parameter #f))
(define *arg1* (make-parameter #f))
(define *arg2* (make-parameter #f))
(define *pc* (make-parameter '()))
;; stack
(define *stack* (make-vector 100))
(define *stack-index* (make-parameter 0))
(define (stack-push v)
  (vector-set! *stack* (*stack-index*) v)
  (*stack-index* (1+ (*stack-index*))))
(define (stack-pop)
  (*stack-index* (1- (*stack-index*)))
  (vector-ref *stack* (*stack-index*)))
;; Leftover bs
(define sg.init (make-vector 100))
(define sg.current (make-vector 100))
(define g.init (make-parameter '()))
(define g.current (make-parameter '()))
(define desc.init (make-parameter '()))

;; Macros
(define-syntax-rule (define-leaf (leaf-name args ...) body ...)
  (begin (define (leaf-name args ...)
           (list (lambda () #((name . leaf-name)) body ...)))
         (set! leafs (cons '(leaf-name (args ...)) leafs))))
(define-syntax-rule (define-compound (name args ...) body ...)
  (begin (define (name args ...) body ...)
         (set! compounds (cons '(name (args ...)) compounds))))

(define-syntax with-fresh-globals
  (syntax-rules ()
    [(_ expr ...)
     (parameterize ([g.current '()]
                    [g.init '()]
                    [desc.init '()]
                    [*env* #f]
                    [*pc* '()]
                    [*stack-index* 0])
       expr ...)]))
(define-syntax define-initial
  (syntax-rules ()
    [(_ name value)
     (vector-set! sg.init (g.init-extend! 'name) value)]))
(define-syntax define-primitive
  (syntax-rules ()
    [(_ name value 0) (define-primitive0 name value)]
    [(_ name value 1) (define-primitive1 name value)]
    [(_ name value 2) (define-primitive2 name value)]
    [(_ name value 3) (define-primitive3 name value)]))
(define-syntax-rule (define-primitive0 name value)
  (define-initial name
    (description-extend! 'name `(function ,value))))
(define-syntax-rule (define-primitive1 name value)
  (define-initial name
    (description-extend! 'name `(function ,value a))))
(define-syntax-rule (define-primitive2 name value)
  (define-initial name
    (description-extend! 'name `(function ,value a b))))
(define-syntax-rule (define-primitive3 name value)
  (define-initial name
    (description-extend! 'name `(function ,value a b c))))

;; Record types
(define-record-type <activation-record>
  (make-activation-record next vals)
  activation-record?
  (next activation-record-next activation-record-next-set!)
  (vals activation-record-vals))
(define-record-type <closure>
  (make-closure code closed-environment)
  closure?
  (code closure-code)
  (closed-environment closure-closed-environment))

;; Helpers
(define (find-indexed pred lst)
  (find (lambda (kons) (pred (car kons)))
        (zip lst (iota (length lst)))))

(define (local-variable env ref)
  (let scan-env ([env env] [i 0])
    (if (null? env)
        #f
        (let ([the-match (find-indexed (lambda (name) (eq? name ref)) (car env))])
          (if the-match
              `(local ,i . ,(cadr the-match))
              (scan-env (cdr env) (+ i 1)))))))

(define (global-variable env ref)
  (assq-ref env ref))

(define (compute-kind ref r)
  (or (local-variable r ref)
      (global-variable (g.current) ref)
      (global-variable (g.init) ref)))

(define (global-update! i v)
  (vector-set! sg.current i v))

(define (g.current-extend! ref)
  (let ([level (length (g.current))])
    (g.current (cons (cons ref (cons 'global level)) (g.current)))
    level))

(define (g.init-extend! ref)
  (let ([level (length (g.init))])
    (g.init (cons (cons ref (cons 'predefined level)) (g.init)))
    level))

(define (static-error . args)
  (apply error "Static Error:" args))

;; activation records
(define (make-activation-rec size)
  (make-activation-record #f (make-vector size)))

(define activation-rec-set!
  (case-lambda
    [(r k v) (vector-set! (activation-record-vals r) k v)]
    [(r j k v) (if (zero? j)
                   (activation-rec-set! r k v)
                   (activation-rec-set! (activation-rec-next r) (1- j) k v))]))

(define activation-rec-ref
  (case-lambda
    [(r i)
     (vector-ref (activation-record-vals r) i)]
    [(r i j)
     (if (zero? i)
         (activation-rec-ref r j)
         (activation-rec-ref (activation-rec-next r) (1- i) j))]))

(define (activation-rec-length r)
  (vector-length (activation-record-vals r)))

(define (activation-rec-next r)
  (if r
      (activation-record-next r)
      (error "Tried to take next of nil record")))

(define (listify! v* arity)
  (let lp ([index (1- (activation-rec-length v*))]
           [result '()])
    (if (= arity index)
        (activation-rec-set! v* arity result)
        (lp (1- index) (cons (activation-rec-ref v* (1- index))
                             result)))))
;; closures
(define (invoke f)
  (cond
   [(closure? f)
    (stack-push (*pc*))
    (*env* (closure-closed-environment f))
    (*pc* (closure-code f))]
   [else (error "Tried to invoke a non-closure")]))

;; function descriptions
(define (description-extend! name description)
  (desc.init (cons (cons name description) (desc.init))))

(define (get-description name)
  (let ([p (assq name (desc.init))])
    (and (pair? p) (cdr p))))

(define desc-args cddr)

(define desc-address cadr)

(define (list->activation-rec lst)
  (let lp ([lst lst]
           [next #f])
    (if (null? lst)
        next
        (lp (cdr lst)
            (make-activation-record next (list->vector (append (car lst) (list undefined-value))))))))

(define (env-set! lst)
  (*env* (list->activation-rec lst)))

;; environments
(define (predefined-fetch i)
  (vector-ref sg.init i))

(define (global-fetch i)
  (vector-ref sg.current i))

(define (r-extend* r n*)
  (cons n* r))

(define (sr-extend* sr v*)
  (activation-record-next-set! v* sr)
  v*)

;; Initial definitions
(define-initial t #t)
(define-initial f #f)
(define-initial nil '())
(define-primitive cons cons 2)
(define-primitive car car 1)
(define-primitive cdr cdr 1)
(define-primitive + + 2)
(define-primitive - - 2)
(define-primitive = = 2)
(define-primitive < < 2)

;; Combinators
(define-leaf (CONSTANT v)
  (*val* v))

(define-leaf (PUSH-VALUE)
  (stack-push (*val*)))

(define-leaf (POP-FUNCTION)
  (*fun* (stack-pop)))

(define-leaf (PRESERVE-ENV)
  (stack-push (*env*)))

(define-leaf (FUNCTION-INVOKE)
  (invoke (*fun*)))

(define-leaf (RESTORE-ENV)
  (*env* (stack-pop)))

(define-leaf (RETURN)
  (*pc* (stack-pop)))

(define-compound (SEQUENCE m m+)
  (append m m+))

(define-leaf (SHALLOW-ARGUMENT-REF j)
  (*val* (activation-rec-ref (*env*) j)))

(define-leaf (DEEP-ARGUMENT-REF i j)
  (*val* (activation-rec-ref (*env*) i j)))

(define-leaf (PREDEFINED i)
  (*val* (predefined-fetch i)))

(define-leaf (GLOBAL-REF i)
  (*val* (global-fetch i)))

(define-leaf (CHECKED-GLOBAL-REF i)
  (let ([v (global-fetch i)])
    (if (eq? v undefined-value)
        (error "Uninitialized variable")
        (*val* v))))

(define-leaf (JUMP-FALSE i)
  (when (not (*val*))
    (*pc* (list-tail (*pc*) i))))

(define-leaf (GOTO i)
  (*pc* (list-tail (*pc*) i)))

(define-compound (ALTERNATIVE m1 m2 m3)
  (append m1 (JUMP-FALSE (1+ (length m2))) m2 (GOTO (length m3)) m3))

(define-compound (STORE-ARGUMENT m m* rank)
  (append m (PUSH-VALUE) m* (POP-FRAME! rank)))

(define-compound (CONS-ARGUMENT m m* arity)
  (append m (PUSH-VALUE) m* (POP-CONS-FRAME! arity)))

(define-leaf (POP-FRAME! rank)
  (activation-rec-set! (*val*) rank (stack-pop)))

(define-leaf (POP-CONS-FRAME! arity)
  (activation-rec-set!
   (*val*) arity
   (cons (stack-pop) (activation-rec-ref (*val*) arity))))

(define-leaf (ALLOCATE-FRAME size)
  (*val* (make-activation-rec (1+ size))))

(define-leaf (ALLOCATE-DOTTED-FRAME size)
  (let ([rec (make-activation-rec (1+ size))])
    (activation-rec-set! rec size '())
    (*val* rec)))

(define-compound (REGULAR-CALL m m+)
  (append
   m
   (PUSH-VALUE)
   m+
   (POP-FUNCTION)
   (PRESERVE-ENV)
   (FUNCTION-INVOKE)
   (RESTORE-ENV)))

(define-compound (TR-REGULAR-CALL m m*)
  (append
   m
   (PUSH-VALUE)
   m*
   (POP-FUNCTION)
   (FUNCTION-INVOKE)))

(define-compound (CALL0 address)
  (INVOKE0 address))

(define-compound (CALL1 address m1)
  (append m1 (INVOKE1 address)))

(define-compound (CALL2 address m1 m2)
  (append m1 (PUSH-VALUE) m2 (POP-ARG1) (INVOKE2 address)))

(define-compound (CALL3 address m1 m2 m3)
  (append m1 (PUSH-VALUE) m2 (PUSH-VALUE) m3 (POP-ARG2) (POP-ARG1) (INVOKE3 address)))

;; Predefined invokers, will eventually be specialized

(define-leaf (INVOKE0 address)
  (*val* (address)))

(define-leaf (INVOKE1 address)
  (*val* (address (*val*))))

(define-leaf (INVOKE2 address)
  (*val* (address (*val*) (*arg1*))))

(define-leaf (INVOKE3 address)
  (*val* (address (*val*) (*arg1*) (*arg2*))))

(define-leaf (POP-ARG1)
  (*arg1* (stack-pop)))

(define-leaf (POP-ARG2)
  (*arg2* (stack-pop)))

(define-compound (FIX-LET m* m+)
  (append m* (EXTEND-ENV) m+ (UNLINK-ENV)))

(define-compound (TR-FIX-LET m* m+)
  (append m* (EXTEND-ENV) m+))

(define-leaf (EXTEND-ENV)
  (*env* (sr-extend* (*env*) (*val*))))

(define-leaf (UNLINK-ENV)
  (*env* (activation-rec-next (*env*))))

(define-compound (SHALLOW-ARGUMENT-SET! j m)
  (append m (SET-SHALLOW-ARGUMENT! j)))

(define-leaf (SET-SHALLOW-ARGUMENT! j)
  (activation-rec-set! (*env*) j (*val*)))

(define-leaf (DEEP-ARGUMENT-SET! i j m)
  (activation-rec-set! (*env*) i j (m)))

(define-compound (GLOBAL-SET! i m)
  (append m (SET-GLOBAL! i)))

(define-leaf (SET-GLOBAL! i)
  (global-update! i (*val*)))

(define-leaf (CREATE-CLOSURE offset)
  (*val* (make-closure (list-tail (*pc*) offset) (*env*))))

(define-leaf (PACK-FRAME! arity)
  (listify! (*val*) arity))

(define-leaf (ARITY>=? arity)
  (unless (>= (activation-rec-length (*val*)) arity)
    (error "Incorrect arity for nary function")))

(define-leaf (ARITY=? arity)
  (unless (= (activation-rec-length (*val*)) arity)
    (error "Incorrect arity for fix function")))

(define-compound (FIX-CLOSURE m+ arity)
  (define the-function
    (append (ARITY=? (1+ arity)) (EXTEND-ENV) m+ (RETURN)))
  (append (CREATE-CLOSURE 1) (GOTO (length the-function)) the-function))

(define-compound (NARY-CLOSURE m+ arity)
  (define the-function
    (append (ARITY>=? (1+ arity)) (PACK-FRAME! arity) (EXTEND-ENV) m+ (RETURN)))
  (append (CREATE-CLOSURE 1) (GOTO (length the-function)) the-function))

(define-leaf (FINISH)
  (*pc* '()))

;; Denotations
(define (meaning e r tail?)
  "Core of the denotation. `tail?' allows us to avoid capturing the local environment if we know we will not return"
  (match e
    [(? (negate pair?))
     (if (symbol? e)      (meaning-reference e r)
         (meaning-quotation e))]
    [('quote e* ...)      (meaning-quotation (car e*))]
    [('lambda v* e* ...)  (meaning-abstraction v* e* r)]
    [('if e1 e2 e3)       (meaning-alternative e1 e2 e3 r tail?)]
    [('begin e* ...)      (meaning-sequence e* r tail?)]
    [('set! n e)          (meaning-assignment n e r)]
    [(f v* ...)           (meaning-application f v* r tail?)]))

(define (meaning-quotation e*)
  (CONSTANT e*))

(define (meaning-alternative e1 e2 e3 r tail?)
  (let ([m1 (meaning e1 r #f)]
        [m2 (meaning e2 r tail?)]
        [m3 (meaning e3 r tail?)])
    (ALTERNATIVE m1 m2 m3)))

(define (meaning-sequence e+ r tail?)
  (cond
   [(null? e+) (static-error "Illegal syntax (begin)")]
   [(null? (cdr e+)) (meaning (car e+) r tail?)]
   [else (SEQUENCE (meaning (car e+) r #f)
                   (meaning-sequence (cdr e+) r tail?))]))

(define (meaning-reference n r)
  (match (compute-kind n r)
    [('local . (i . j))
     (if (= i 0)
         (SHALLOW-ARGUMENT-REF j)
         (DEEP-ARGUMENT-REF i j))]
    [('global . i)
     (CHECKED-GLOBAL-REF i)]
    [('predefined . i)
     (PREDEFINED i)]
    [_ (static-error "No such variable" n)]))

(define (meaning-application f e* r tail?)
  (cond
   [(and (symbol? f)
         (get-description f)
         (compute-kind f r)) =>
         (lambda (kind)
           (if (eq? (car kind) 'predefined)
               (meaning-primitive-application f e* r)
               (meaning-regular-application f e* r tail?)))]
   [(and (pair? f) (eq? (car f) 'lambda))
    (meaning-closed-application f e* r tail?)]
   [else (meaning-regular-application f e* r tail?)]))

(define (meaning-list e* r tail?)
  (cond
   [(null? e*) '()]
   [(null? (cdr e*)) (list (meaning (car e*) r tail?))]
   [else (cons (meaning (car e*) r #f) (meaning-list (cdr e*) r tail?))]))

(define (meaning* e* r)
  (let lp ([rank 0]
           [m* (meaning-list e* r #f)])
    (if (pair? m*)
        (STORE-ARGUMENT (car m*) (lp (1+ rank) (cdr m*)) rank)
        (ALLOCATE-FRAME (length e*)))))

(define (meaning*-dotted e* r fixed-size)
  (let lp ([rank 0]
           [m* (meaning-list e* r #f)])
    (cond
     [(null? m*) (ALLOCATE-DOTTED-FRAME fixed-size)]
     [(> fixed-size rank) (STORE-ARGUMENT (car m*) (lp (1+ rank) (cdr m*)) rank)]
     [else (CONS-ARGUMENT (car m*) (lp (1+ rank) (cdr m*)) fixed-size)])))

(define (meaning-regular-application f e* r tail?)
  (let [(mf (meaning f r #f))
        (m* (meaning* e* r))]
    ((if tail? TR-REGULAR-CALL REGULAR-CALL) mf m*))) ;; hehe new trick

(define (meaning-primitive-application f e* r)
  (let* ([desc (get-description f)]
         [address (desc-address desc)]
         [arity (length e*)])
    (if (= (length (desc-args desc)) arity)
        (case arity
          [(0) (CALL0 address)]
          [(1) (let ([m1 (meaning (first e*) r #f)])
                 (CALL1 address m1))]
          [(2) (let ([m1 (meaning (first e*)  r #f)]
                     [m2 (meaning (second e*) r #f)])
                 (CALL2 address m1 m2))]
          [(3) (let ([m1 (meaning (first e*)  r #f)]
                     [m2 (meaning (second e*) r #f)]
                     [m3 (meaning (third e*)  r #f)])
                 (CALL3 address m1 m2 m3))])
        (static-error "Wrong arity for" f arity))))

(define (meaning-closed-application e ee* r tail?)
  (let lp ([n* (second e)]
           [e* ee*]
           [fixargs '()])
    (cond
     [(pair? n*)
      (if (pair? e*)
          (lp (cdr n*) (cdr e*) (cons (car n*) fixargs))
          (static-error "Not enough arguments" e ee*))]
     [(null? n*)
      (if (pair? e*)
          (static-error "Too many arguments" e ee*)
          (meaning-fixed-closed-application (second e) (drop e 2) ee* r tail?))]
     [else
      (meaning-dotted-closed-application (reverse fixargs) n* (drop e 2) ee* r tail?)])))

(define (meaning-fixed-closed-application n* body e* r tail?)
  (let* [(m* (meaning* e* r))
         (r2 (r-extend* r n*))
         (m+ (meaning-sequence body r2 tail?))]
    ((if tail? TR-FIX-LET FIX-LET) m* m+)))

(define (meaning-dotted-closed-application n* n body e* r tail?)
  (let* ([m* (meaning*-dotted e* r (length n*))]
         [r2 (r-extend* r (append n* (list n)))]
         [m+ (meaning-sequence body r2 tail?)])
    ((if tail? TR-FIX-LET FIX-LET) m* m+)))

(define (meaning-assignment n e r)
  (let ([m (meaning e r #f)])
    (match (compute-kind n r)
      [('local . (i . j))
       (if (zero? i)
           (SHALLOW-ARGUMENT-SET! j m)
           (DEEP-ARGUMENT-SET! i j m))]
      [('global . i)
       (GLOBAL-SET! i m)]
      [('predefined _ ...)
       (static-error "Cannot set predefined variable" n)]
      [_ (static-error "Cannot set! undefined variable" n)])))

(define (meaning-abstraction nn* body r)
  (let lp ([n* nn*]
           [fixargs '()])
    (cond
     [(pair? n*) (lp (cdr n*) (cons (car n*) fixargs))]
     [(null? n*) (meaning-fixed-abstraction nn* body r)]
     [else (meaning-dotted-abstraction (reverse fixargs) n* body r)])))

(define (meaning-fixed-abstraction n* body r)
  (let* ([arity (length n*)]
         [r2 (r-extend* r n*)]
         [m+ (meaning-sequence body r2 #t)])
    (FIX-CLOSURE m+ arity)))

(define (meaning-dotted-abstraction n* n body r)
  (let* ([arity (length n*)]
         [r2 (r-extend* r (append n* (list n)))]
         [m+ (meaning-sequence body r2 #t)])
    (NARY-CLOSURE m+ arity)))

;; VM Runner
(define (run)
  (let ([instruction (car (*pc*))])
    (*pc* (cdr (*pc*)))
    (instruction)
    (if (pair? (*pc*))
        (run)
        (format #f "Execution terminated with *val* = ~a\n" (*val*)))))

;; Reporting
(letrec ([combinator-name (compose symbol->string car)]
         [sort-combinators
          (lambda (combinators)
            (sort combinators (lambda (a b)
                                (string< (combinator-name a) (combinator-name b)))))]
         [n-leafs (length leafs)]
         [n-compounds (length compounds)])
  (format #t "Defined ~a combinators; ~a leaves and ~a compound instructions\n" (+ n-leafs n-compounds) n-leafs n-compounds)
  (format #t "Leaves:\n")
  (for-each (lambda (c) (format #t " (~a ~a)\n" (first c) (second c)))
            (sort-combinators leafs))
    (format #t "Compounds:\n")
  (for-each (lambda (c) (format #t " (~a ~a)\n" (first c) (second c)))
            (sort-combinators compounds)))

;; Tests
(define (run-test prgm)
  (*pc* prgm)
  (let lp ()
    (define instruction (car (*pc*)))
    (*pc* (cdr (*pc*)))
    (instruction)
    (if (null? (*pc*))
        (*val*)
        (lp))))

(test-begin "bytecode-compiler")
;; quotation
(test-group "meaning-quotation"
  (with-fresh-globals
   (test-eq 'apple
     (run-test (meaning-quotation 'apple)))
   (test-eq '(apple pear)
     (run-test (meaning-quotation '(apple pear))))))
;; alternatives
(test-group "meaning-alternative"
  (with-fresh-globals
   (env-set! '((#t apple pear)))
   (test-eq 'apple
     (run-test (meaning-alternative 'p 't 'f '((p t f)) #t)))
   (env-set! '((#f apple pear)))
   (test-eq 'pear
     (run-test (meaning-alternative 'p 't 'f '((p t f)) #t)))))
  ;; sequences
(test-group "meaning-sequence"
  (with-fresh-globals
   (env-set! '((1 2)))
   (test-eq 2
     (run-test (meaning-sequence '(a b) '((a b)) #t)))))
;; references
(test-group "meaning-reference"
  (with-fresh-globals
   (g.current-extend! 'a)
   (vector-set! sg.current 0 10)
   (g.init-extend! 'b)
   (vector-set! sg.init 0 20)
   (env-set! '((0 1)))
   (test-eq 1
     (run-test (meaning-reference 'b '((a b)))))
   (env-set! '((0 1) (2)))
   (test-eq 0
     (run-test (meaning-reference 'b '((a) (b c)))))
   (test-equal 10
     (run-test (meaning-reference 'a '())))
   (test-equal 20
     (run-test (meaning-reference 'b '())))))
;; applications
(test-group "meaning-application"
  (with-fresh-globals
   (desc.init `((+ function ,+ . (x y))))
   (define *-abstraction
     (run-test (FIX-CLOSURE
                (list (lambda ()
                        (*val*
                         (* (activation-rec-ref (*env*) 0)
                            (activation-rec-ref (*env*) 1)))))
       2)))
   (env-set! `((,*-abstraction)))
   (test-eq 9
     (run-test (meaning-regular-application '* '(3 3) '((*)) #t)))
   (test-eq 2
     (run-test (meaning-primitive-application '+ '(1 1) '(()))))
   (test-eq 2 (run-test (meaning-closed-application '(lambda (x) x) '(2) '(()) #t)))
   (test-equal '(2 3) (run-test (meaning-closed-application '(lambda (x . y) y) '(1 2 3) '(()) #t)))))
;; assignment
(test-group "meaning-assignment"
  (with-fresh-globals
   (*env* (list->activation-rec '((1))))
  (run-test (meaning-assignment 'apples 2 '((apples))))
   (test-eq 2                                
     (activation-rec-ref (*env*) 0 0))))
;; abstraction
(test-group "meaning-abstraction"
  (define test-code (append
                     (PUSH-VALUE)
                     (CONSTANT 1)
                     (PUSH-VALUE)
                     (CONSTANT 2)
                     (PUSH-VALUE)
                     (ALLOCATE-FRAME 2)
                     (POP-FRAME! 1)
                     (POP-FRAME! 0)
                     (POP-FUNCTION)
                     (FUNCTION-INVOKE)))
  (with-fresh-globals
   (test-assert (closure? (run-test (meaning-abstraction '(x) '(x) '(())))))
   (test-eq 2
     (run-test (append (meaning-abstraction '(x y) '(y) '(())) test-code)))
   (test-equal '(2)
     (run-test (append (meaning-abstraction '(x . y) '(y) '(())) test-code)))))
;; (test-group "call/cc"
;;   (test-eq 3 ((meaning '(+ 2 (call/cc (lambda (cc) (cc 1) 2))) '(()) #t)))
;;   (test-eq 5 ((meaning '((lambda (cont v)
;;                            (set! v (+ 2 (call/cc (lambda (cc) (set! cont cc) 1))))
;;                            (if (= v 3)
;;                                (cont v)
;;                                v)) #f 0)
;;                        '(()) #t))))
;; (test-group "apply"
;;   (test-eq 3 ((meaning '(apply + 1 2 '()) '(()) #t)))
;;   (test-eq 3 ((meaning '(apply + '(1 2)) '(()) #t))))
(test-end "bytecode-compiler")