#lang typed/racket
(require "ast.rkt")
(require racket/treelist)
(: prelude CoreProgram)
(define prelude
  '((define (I x) x) (define (K x y) x) (define (K1 x y) y) (define (S f g x) ((f x) (g x))) (define (compose f g x) (f (g x))) (define (twice f) ((compose f) f))))

(struct State ([stack : Stack] [dump : Dump] [heap : Heap] [globals : Globals] [stats : Stats]))
(define-type Addr Index)
(define-type Stack (Listof Addr))
(define-type Dump Null)
(define-type Heap (TreeListof Node))
(define-type Node (∪ (List Addr Addr) CoreScDefn Integer (List 'Indirect Addr) 'Undef))
(define-type Globals (Immutable-HashTable Name Addr))
(define-type Stats Nonnegative-Integer)


(: ap-node? (-> Node Boolean : (List Addr Addr)))
(define (ap-node? a)
  (and (pair? a) (number? (car a)) (pair? (cdr a)) (null? (cddr a))))

(: undef-node? (-> Any Boolean : 'Undef))
(define (undef-node? a)
  (eqv? a 'Undef))

(: indirect-node? (-> Node Boolean : (List 'Indirect Addr)))
(define (indirect-node? a)
  (and (pair? a) (eqv? (car a) 'Indirect)))

(: indirect-node (-> Addr (∩ (List 'Indirect Addr) Node)))
(define (indirect-node a)
  (list 'Indirect a))

(: integer-node? (-> Node Boolean : Integer))
(define (integer-node? a)
  (integer? a))

(: sc-node? (-> Node Boolean : CoreScDefn))
(define (sc-node? a)
  (and (pair? a) (eqv? (car a) 'define)))

(: ap-fun (-> (List Addr Addr) Addr))
(define (ap-fun a)
  (first a))

(: ap-arg (-> (List Addr Addr) Addr))
(define (ap-arg a)
  (second a))

(: match-sc (-> CoreScDefn (Values (Listof Name) CoreExpr)))
(define (match-sc scdef)
  (values (cdadr scdef) (caddr scdef)))

(: indirect-addr (-> (List 'Indirect Addr) Addr))
(define (indirect-addr a)
  (second a))

(: empty-heap Heap)
(define empty-heap
  (treelist))

(: empty-globals Globals)
(define empty-globals
  (make-immutable-hash))

(: empty-dump Dump)
(define empty-dump '())

(: empty-stats Stats)
(define empty-stats 0)

(: empty-stack Stack)
(define empty-stack '())

(: push-addr (-> Stack Addr (Pairof Addr Stack)))
(define (push-addr st n)
  (cons n st))

(: incr-stats (-> Stats Stats))
(define incr-stats add1)

(: update-heap-undef (-> Heap Addr Node Heap))
(define (update-heap-undef heap addr node)
  (if (eq? (treelist-ref heap addr) 'Undef)
      (update-heap heap addr node)
      (error "Address in heap is not undefined")))

(: update-heap (-> Heap Addr Node Heap))
(define (update-heap heap addr node)
  (treelist-set heap addr node))


(: update-globals (-> Globals Name Addr Globals))
(define (update-globals globals name addr)
  (hash-set globals name addr))

(: update-stats (-> (-> Stats Stats) State State))
(define (update-stats f mstate)
  (struct-copy State mstate [stats (f (State-stats mstate))]))

(: update-stack (-> (-> Stack Stack) State State))
(define (update-stack f mstate)
  (struct-copy State mstate [stack (f (State-stack mstate))]))

(: allocate-node (-> Heap Node (Values Heap Addr)))
(define (allocate-node heap node)
  (let ([addr (treelist-length heap)]
        [heap (treelist-add heap node)])
    (values heap addr)))

(: lookup-node (-> Heap Addr Node))
(define lookup-node treelist-ref)

(: lookup-globals (-> Globals Name Addr))
(define lookup-globals hash-ref)

(: initialize-heap (-> (Listof CoreScDefn) (Values Heap Globals)))
(define (initialize-heap scs)
  (for/fold ([heap : Heap empty-heap]
             [globals : Globals empty-globals])
            ([sc : CoreScDefn scs])
    (let-values ([(new-heap addr) (allocate-node heap sc)])
      (values new-heap
              (update-globals globals (scdefn-name sc) addr)))))

(: compile-core (-> (Listof CoreScDefn) State))
(define (compile-core scdefs)
  (let-values ([(heap globals) (initialize-heap (append scdefs prelude))])
    (State (list (lookup-globals globals 'main)) empty-dump heap globals empty-stats)))


(: eval-state (-> State (Pairof State (Listof State))))
(define (eval-state st)
  (if (final-state? st)
      (list st)
      (let ([next-st (update-stats incr-stats (step st))])
        (cons st (eval-state next-st)))))

(: final-state? (-> State Boolean))
(define (final-state? st)
  (match (State-stack st)
    [(list n) (integer-node? (lookup-node (State-heap st) n))]
    [_ #f]))

(: step (-> State State))
(define (step state)
  (let ([stack (State-stack state)])
    (match stack
      [(cons addr rstack)
       (step-node addr rstack state (lookup-node (State-heap state) addr))]
      ['() (error "Invalid state")])))

(: step-node (-> Addr Stack State Node State))
(define (step-node addr rstack st n)
  (cond
    [(integer-node? n) (error "A number cannot be stepped any further!")]
    [(ap-node? n) (update-stack (λ (stack) (push-addr stack (ap-fun n))) st)]
    [(undef-node? n) (error "An undefined node is being stepped!")]
    [(sc-node? n)
     (let-values ([(args body) (match-sc n)])
       (step-sc addr rstack st args body))]
    [(indirect-node? n)
     (update-stack (λ (_) (push-addr rstack (indirect-addr n))) st)]))

(: get-arg (-> Heap Addr Addr))
(define (get-arg heap addr)
  (match (lookup-node heap addr)
    [(list _ a) a]
    [_ (error "Not an application node")]))

(: step-sc (-> Addr Stack State (Listof Name) CoreExpr State))
(define (step-sc root rstack state arg-names body)
  (let ([heap (State-heap state)])
    (let*-values ([(new-env new-stack root)
                  (for/fold ([globals : Globals (State-globals state)]
                             [stack : Stack rstack]
                             [root : Addr root])
                            ([arg-name : Name arg-names])
                    (if (null? stack)
                        (error "Not enough arguments to apply the supercombinator")
                        (values (update-globals globals arg-name (get-arg heap (car stack))) (cdr stack) (car stack))))])
      (let* ([heap (instantiate-and-assign body new-env heap root)]
             [new-stack (cons root new-stack)])
        (struct-copy State state [stack new-stack] [heap heap])))))

(: instantiate-and-allocate (-> CoreExpr Globals Heap (Values Heap Addr)))
(define (instantiate-and-allocate e env heap)
  (let-values ([(heap addr-or-node) (instantiate-expr e env heap)])
    (cond
      [(allocated-addr? addr-or-node)
       (values heap addr-or-node)]
      [else
       (allocate-node heap (unallocated-node-get addr-or-node))])))

(: instantiate-and-assign (-> CoreExpr Globals Heap Addr Heap))
(define (instantiate-and-assign e env heap addr)
  (let-values ([(heap addr-or-node) (instantiate-expr e env heap)])
    (cond
      [(allocated-addr? addr-or-node)
       (update-heap heap addr (indirect-node addr-or-node))]
      [else
       (update-heap heap addr (unallocated-node-get addr-or-node))])))

(: instantiate-expr (-> CoreExpr Globals Heap (Values Heap (∪ Addr Unallocated-Node))))
(define (instantiate-expr e env heap)
  (match e
    [(? integer?) (values heap (unallocated-node e))]
    [(list f a)
     (let*-values ([(heap addr-f) (instantiate-and-allocate f env heap)]
                   [(heap addr-a) (instantiate-and-allocate a env heap)])
       (values heap (unallocated-node (list addr-f addr-a))))]
    [(? symbol?) (values heap (lookup-globals env e))]
    [(list 'let binds e)
     (let-values ([(env heap) (instantiate-binds binds env heap)])
       (instantiate-expr e env heap))]
    [(list 'letrec binds e)
     (let-values ([(env heap) (instantiate-rec-binds binds env heap)])
       (instantiate-expr e env heap))]
    [_ (error "unimplemented")]))

(: instantiate-rec-binds (-> CoreBinds Globals Heap (Values Globals Heap)))
(define (instantiate-rec-binds binds env heap)
  (let-values ([(addrs env heap)
                (for/fold
                    ([addrs : (Listof Addr) '()]
                     [env : Globals env]
                     [heap : Heap heap])
                    ([bind : CoreBind binds])
                  (let-values ([(heap addr) (allocate-node heap 'Undef)])
                    (values (cons addr addrs) (update-globals env (first bind) addr) heap)))])
    (values env (for/fold
        ([heap : Heap heap])
        ([bind : CoreBind binds]
         [addr : Addr (reverse addrs)])
                  (instantiate-rec-bind addr bind env heap)))))

(: instantiate-rec-bind (-> Addr CoreBind Globals Heap Heap))
(define (instantiate-rec-bind addr bind env heap)
  (let ([x (first bind)]
        [e (second bind)])
    (let ([heap (instantiate-and-assign e env heap addr)])
      heap)))

(: instantiate-binds (-> CoreBinds Globals Heap (Values Globals Heap)))
(define (instantiate-binds binds old-env heap)
  (for/fold
      ([env : Globals old-env]
       [heap : Heap heap])
      ([bind : CoreBind binds])
    (instantiate-bind bind old-env env heap)))

(: instantiate-bind (-> CoreBind Globals Globals Heap (Values Globals Heap)))
(define (instantiate-bind bind old-env new-env heap)
  (let ([x (first bind)]
        [e (second bind)])
    (let*-values ([(heap addr) (instantiate-and-allocate e old-env heap)])
      (values (update-globals new-env x addr) heap))))


(define-type Unallocated-Node (Pair 'Unallocated Node))

(: unallocated-node-get (-> Unallocated-Node Node))
(define (unallocated-node-get a)
  (cdr a))

(: unallocated-node? (-> (∪ Addr Unallocated-Node) Boolean : Unallocated-Node))
(define (unallocated-node? a)
  (pair? a))

(: allocated-addr? (-> (∪ Addr Unallocated-Node) Boolean : Addr))
(define (allocated-addr? a)
  (number? a))

(: unallocated-node (-> Node Unallocated-Node))
(define (unallocated-node a)
  (cons 'Unallocated a))

(provide (all-defined-out))