(use-modules (ice-9 match)
             (srfi srfi-1) 
             (srfi srfi-9) ;; Record types
             (srfi srfi-64) ;; Testing framework
             )
;; This final interpreter removes reified continuations, favoring a magic operator inherited
;; from the implementing language. This allows us to represent procedures as thunks, and
;; the procedures that create them as combinators
;; (note: is this true? don't they still use free variables in their definitions?)
;; ((yes, because they are of dynamic scope, therefore not captured))

;; Globals
(define undefined-value (make-symbol "undefined"))
(define g.init (make-parameter '()))
(define g.current (make-parameter '()))
(define desc.init (make-parameter '()))
(define *env* (make-parameter #f))

(define sg.init (make-vector 100))
(define sg.current (make-vector 100))

;; Macros
(define-syntax with-fresh-globals
  (syntax-rules ()
    [(_ expr ...)
     (parameterize ([g.current '()]
                    [g.init '()]
                    [desc.init '()]
                    [*env* #f])
       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
    (letrec ([arity+1 (+ 0 1)]
             [behavior
              (lambda (v* sr)
                (if (= (activation-rec-length v*) arity+1)
                    (value)
                    (error "Incorrect arity" 'name)))])
      (description-extend! 'name `(function ,value))
      (make-closure behavior (*env*)))))
(define-syntax-rule (define-primitive1 name value)
  (define-initial name
    (letrec ([arity+1 (+ 1 1)]
             [behavior
              (lambda (v* k)
                (if (= (activation-rec-length v*) arity+1)
                    (k (value (activation-rec-ref v* 0)))
                    (error "Incorrect arity" 'name)))])
      (description-extend! 'name `(function ,value a))
      (make-closure behavior (*env*)))))
(define-syntax-rule (define-primitive2 name value)
  (define-initial name
    (letrec ([arity+1 (+ 2 1)]
             [behavior
              (lambda (v* sr)
                (if (= (activation-rec-length v*) arity+1)
                    (value (activation-rec-ref v* 0)
                           (activation-rec-ref v* 1))
                    (error "Incorrect arity" 'name)))])
      (description-extend! 'name `(function ,value a b))
      (make-closure behavior (*env*)))))
(define-syntax-rule (define-primitive3 name value)
  (define-initial name
    (letrec ([arity+1 (+ 3 1)]
             [behavior
              (lambda (v* sr)
                (if (= (activation-rec-length v*) arity+1)
                    (value (activation-rec-ref v* 0)
                           (activation-rec-ref v* 1)
                           (activation-rec-ref v* 2))
                    (error "Incorrect arity" 'name)))])
      (description-extend! 'name `(function ,value a b c))
      (make-closure behavior (*env*)))))

;; 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-record-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-record-next r) (1- i) j))]))
(define (activation-rec-length r)
  (vector-length (activation-record-vals r)))

(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 v*)
  (if (closure? f)
      ((closure-code f) v* (closure-closed-environment f))
      (error "Not a function" f)))
;; 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-initial call/cc
  (make-closure
   (lambda (v* sr)
     (if (= (activation-rec-length v*) 2)
         (call/cc ;; call/cc is "magic" for our purposes here,
                  ;; there are no reified continuations in the denotation
          (lambda (k)
            (invoke
             (activation-rec-ref v* 0)
             (let ([rec (make-activation-rec 2)])
               (activation-rec-set!
                rec 0
                (make-closure
                 (lambda (values r)
                   (if (= (activation-rec-length values) 2)
                       (k (activation-rec-ref values 0))
                       (error "Incorrect arity" 'continuation)))
                 (*env*)))
               rec))))
         (error "Incorrect arity" 'call/cc)))
   (*env*)))
(define-initial apply
  (make-closure
   (lambda (v* sr)
     (if (>= (activation-rec-length v*) 3)
         (let* ([proc (activation-rec-ref v* 0)]
                [last-arg-index (- (activation-rec-length v*) 2)]
                [last-arg (activation-rec-ref v* last-arg-index)]
                [size (+ last-arg-index (length last-arg))]
                [rec (make-activation-rec size)])
           (do ([i 1 (+ i 1)])
               ((= i last-arg-index))
             (activation-rec-set! rec (- i 1) (activation-rec-ref v* i)))
           (do ([i (- last-arg-index 1) (+ i 1)]
                [last-arg last-arg (cdr last-arg)])
               ((null? last-arg))
             (activation-rec-set! rec i (car last-arg)))
           (invoke proc rec))
         (error "Wrong arity" 'apply)))
   (*env*)))

;; Combinators
(define (CONSTANT v)
  (lambda () v))

(define (SEQUENCE m m+)
  (lambda () (m) (m+)))

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

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

(define (PREDEFINED i)
  (lambda () (predefined-fetch i)))

(define (GLOBAL-REF i)
  (lambda () (global-fetch i)))

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

(define (ALTERNATIVE m1 m2 m3)
  (lambda () (if (m1) (m2) (m3))))

(define (STORE-ARGUMENT m m* rank)
  (lambda ()
    (let ([v* (m*)]
          [v (m)])
      (activation-rec-set! v* rank v)
      v*)))

(define (CONS-ARGUMENT m m* size)
  (lambda ()
    (let* ([v* (m*)]
           [v (m)]
           [lst (activation-rec-ref v* size)])
      (activation-rec-set! v* size (cons v lst))
      v*)))

(define (ALLOCATE-FRAME size)
  (lambda () (make-activation-rec (1+ size))))

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

(define (REGULAR-CALL m m*)
  (lambda ()
    (let* ([f (m)]
           [v* (m*)]
           [sr (*env*)]
           [res (invoke f v*)])
      (*env* sr)
      res)))

(define (TR-REGULAR-CALL m m*)
  (lambda () (invoke (m) (m*))))

(define (CALL0 address)
  (lambda () (address)))

(define (CALL1 address m1)
  (lambda () (address (m1))))

(define (CALL2 address m1 m2)
  (lambda ()
    (let ([v1 (m1)])
      (address v1 (m2)))))

(define (CALL3 address m1 m2 m3)
  (lambda ()
    (let ([v1 (m1)]
          [v2 (m2)])
      (address v1 v2 (m3)))))

(define (FIX-LET m* m+)
  (lambda ()
    (*env* (sr-extend* (*env*) (m*)))
    (let ([res (m+)])
      (*env* (activation-record-next (*env*))) ;; black magic
      res)))

(define (TR-FIX-LET m* m+)
  (lambda ()
    (*env* (sr-extend* (*env*) (m*)))
    (m+)))

(define (SHALLOW-ARGUMENT-SET! j m)
  (lambda () (activation-rec-set! (*env*) j (m))))

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

(define (GLOBAL-SET! i m)
  (lambda () (global-update! i (m))))

(define (FIX-CLOSURE m+ arity)
  (lambda ()
    (define (the-function v* sr)
      (if (= (activation-rec-length v*) (1+ arity))
          (begin (*env* (sr-extend* sr v*))
                 (m+))
          (error "Incorrect arity")))
    (make-closure the-function (*env*))))

(define (NARY-CLOSURE m+ arity)
  (lambda ()
    (define (the-function v* sr)
      (if (>= (activation-rec-length v*) (1+ arity))
          (begin (listify! v* arity)
                 (*env* (sr-extend* sr v*))
                 (m+))
          (error "Incorrect arity")))
    (make-closure the-function (*env*))))

;; 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)))

;; Tests
(test-begin "direct-style-interpreter")
;; quotation
(test-group "meaning-quotation"
  (test-eq 'apple
    ((meaning-quotation 'apple)))
  (test-eq '(apple pear)
    ((meaning-quotation '(apple pear)))))
;; alternatives
(test-group "meaning-alternative"
  (with-fresh-globals
   (env-set! '((#t apple pear)))
   (test-eq 'apple
     ((meaning-alternative 'p 't 'f '((p t f)) #t)))
   (env-set! '((#f apple pear)))
   (test-eq 'pear
     ((meaning-alternative 'p 't 'f '((p t f)) #t)))))
  ;; sequences
(test-group "meaning-sequence"
  (with-fresh-globals
   (env-set! '((1 2)))
   (test-eq 2
     ((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
     ((meaning-reference 'b '((a b)))))
   (env-set! '((0 1) (2)))
   (test-eq 0
     ((meaning-reference 'b '((a) (b c)))))
   (test-equal 10
     ((meaning-reference 'a '())))
   (test-equal 20
     ((meaning-reference 'b '())))))
;; applications
(test-group "meaning-application"
  (with-fresh-globals
   (desc.init `((+ function ,+ . (x y))))
   (define *-abstraction
     (FIX-CLOSURE
      (lambda ()
        (* (activation-rec-ref (*env*) 0)
           (activation-rec-ref (*env*) 1)))
      2))
   (env-set! `((,(*-abstraction))))
   (test-eq 9
     ((meaning-regular-application '* '(3 3) '((*)) #t)))
   (test-eq 2
     ((meaning-primitive-application '+ '(1 1) '(()))))
   (test-eq 2 ((meaning-closed-application '(lambda (x) x) '(2) '(()) #t)))
   (test-equal '(2 3) ((meaning-closed-application '(lambda (x . y) y) '(1 2 3) '(()) #t)))))
;; assignment
(test-group "meaning-assignment"
  (with-fresh-globals
   (*env* (list->activation-rec '((1))))
   ((meaning-assignment 'apples 2 '((apples))))
   (test-eq 2                                
     (activation-rec-ref (*env*) 0 0))))
;; abstraction
(test-group "meaning-abstraction"
  (with-fresh-globals
   (test-assert (closure? ((meaning-abstraction '(x) '(x) '(())))))
   (*env* (sr-extend* (*env*) (make-activation-record #f (vector 1 2 undefined-value))))
   (define (call-with-1-2 f)
     (invoke f (make-activation-record #f (vector 1 2 undefined-value))))
   (test-eq 2
     (call-with-1-2 ((meaning-abstraction '(x y) '(y) '(())))))
   (test-equal '(2)
     (call-with-1-2 ((meaning-abstraction '(x . y) '(y) '(())))))))
(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 "direct-style-interpreter")