#lang plai/gc2/collector
;; metadata size
(define METADATA-SIZE 2)

;; where the start of the current "to space" is stored in the heap.
(define LOC:OFFSET 0)

;; where the allocation pointer is stored in the heap
(define LOC:PTR 1)

;; Start of the current "to space"
(define (off) (heap-ref LOC:OFFSET))

;; Offset into the current "to space" where free space starts.
(define (ptr) (heap-ref LOC:PTR))

;; How big are the semi-spaces?
(define (space-size)
  (quotient (- (heap-size) METADATA-SIZE) 2))

;; All functions named gc:*, along with init-allocator, must be
;; implemented by any plai/gc2 collector. For their functionality and
;; interface see [1].

(define (init-allocator)
  (heap-set! LOC:OFFSET METADATA-SIZE)
  (heap-set! LOC:PTR 0))

(module+ test
  (test (with-heap (make-vector 1000)  (+ METADATA-SIZE (space-size) (space-size))) 1000)
  (test (with-heap (make-vector 999)  (+ METADATA-SIZE (space-size) (space-size))) 998))


;; Define some syntax rules to make it easier to write tests

;; Test if the last two expressions are equal.
;; Takes a vector for a heap

(define-syntax (test/heap stx)
  (syntax-case stx ()
    [(test/heap heap oper ...  expected)
     (syntax-protect
      #`(with-heap heap
          (init-allocator)
          #,(syntax/loc stx
              (test (begin oper ...) expected))))]))

;; Test if one of the expressions before the last throws an exception
;; matching the last expression (a string).
;; Takes a vector for a heap
(define-syntax (test/heap/exn stx)
  (syntax-case stx ()
    [(test/heap heap oper ...  expected)
     (syntax-protect
      #`(with-heap heap
          (init-allocator)
          #,(syntax/loc stx
              (test/exn (begin oper ...) expected))))]))

(define (swap-spaces!)
  ;; do nothing
  (void))

(module+ test

  ;; Initially, allocations are in the left half.
  (test/heap
      (make-vector (+ 4 METADATA-SIZE) '?)
    (gc:alloc-flat #f)
    (current-heap)
    #(2 2 flat #f ? ?))

  ;; After calling swap-spaces!, allocations are in the right half
  (test/heap
      (make-vector (+ 4 METADATA-SIZE) '?)
    (swap-spaces!)
    (gc:alloc-flat #f)
    (current-heap)
    #(4 2 ? ? flat #f))

  ;; Swapping twice is back to allocating in left
  (test/heap
      (make-vector (+ 4 METADATA-SIZE) '?)
    (swap-spaces!)
    (swap-spaces!)
    (gc:alloc-flat #f)
    (current-heap)
    #(2 2 flat #f ? ?))
  )

;; malloc : size -> address
(define (malloc n)
  (when (> (+ (ptr) n) (space-size))
    (gc!))
  (when (> (+ (ptr) n) (space-size))
    (error 'malloc "out of memory!"))
  (heap-set! LOC:PTR (+ (ptr) n))
  (+ (heap-ref LOC:OFFSET) (- (ptr) n)))

(define (gc!)
  ;; do nothing
  (void))

;; gc:alloc-flat : flat-value -> address
(define (gc:alloc-flat value)
  (define addr (malloc 2))
  (heap-set! addr 'flat)
  (heap-set! (+ addr 1) value)
  addr)

;; gc:flat? : address -> boolean
(define (gc:flat? address)
  (equal? (heap-ref address) 'flat))

;; gc:deref : address -> flat-value
(define (gc:deref address)
  (unless (gc:flat? address)
    (error 'gc:deref "not a flat: ~a" address))
  (heap-ref (+ address 1)))

;; gc:cons : root root -> address
(define (gc:cons root1 root2)
  (define addr (malloc 3 root1 root2))
  (heap-set! addr 'cons)
  (heap-set! (+ addr 1) (read-root root1))
  (heap-set! (+ addr 2) (read-root root2))
  addr)

;; gc:cons? : address -> boolean
(define (gc:cons? address)
  (equal? (heap-ref address) 'cons))

;; gc:first : address -> address
(define (gc:first address)
  (unless (gc:cons? address)
    (error 'gc:first "not a pair: ~a" address))
  (heap-ref (+ address 1)))

;; gc:rest : address -> address
(define (gc:rest address)
  (unless (gc:cons? address)
    (error 'gc:rest "not a pair: ~a" address))
  (heap-ref (+ address 2)))

;; gc:set-first! : address address -> void
(define (gc:set-first! address new-value-address)
  (unless (gc:cons? address)
    (error 'gc:set-first! "not a pair: ~a" address))
  (heap-set! (+ address 1) new-value-address))

;; gc:set-rest! : address address -> void
(define (gc:set-rest! address new-value-address)
  (unless (gc:cons? address)
    (error 'gc:set-rest! "not a pair: ~a" address))
  (heap-set! (+ address 2) new-value-address))

;; gc:closure : opaque-value (listof root) -> address
(define (gc:closure code-ptr free-vars)
  (define n-vars (length free-vars))
  (define addr (malloc (+ 3 n-vars)))
  (heap-set! addr 'clos)
  (heap-set! (+ addr 1) code-ptr)
  (heap-set! (+ addr 2) n-vars)
  (for ([i  (in-range n-vars)]
        [fv (in-list free-vars)])
    (heap-set! (+ addr 3 i) (read-root fv)))
  addr)

;; gc:closure? : address -> boolean
(define (gc:closure? address)
  (equal? (heap-ref address) 'clos))

;; gc:closure-code-ptr : address -> opaque-value
(define (gc:closure-code-ptr address)
  (unless (gc:closure? address)
    (error 'gc:closure-code-ptr "not a closure: ~a" address))
  (heap-ref (+ address 1)))

;; gc:closure-env-ref : address integer -> address
(define (gc:closure-env-ref address i)
  (unless (gc:closure? address)
    (error 'gc:closure-env-ref "not a closure: ~a" address))
  (heap-ref (+ address 3 i)))

(module+ test
  ;; OOM
  (test/heap/exn (make-vector METADATA-SIZE)
    (gc:alloc-flat #f)
    "out of memory")

  ;; OOM due to using only half of the heap
  (test/heap/exn
      (make-vector (+ 2 METADATA-SIZE))
    (gc:alloc-flat #f)
    "out of memory")

  ;; dereferencing cons as flat
  (test/heap/exn (make-vector 1000)
    (let ([cons-addr
           (gc:cons
            (simple-root (gc:alloc-flat #f))
            (simple-root (gc:alloc-flat #t)))])
      (gc:deref cons-addr))
    "not a flat")

  ;; dereferencing flat as cons
  (test/heap/exn (make-vector 1000)
    (let ([flat-addr (gc:alloc-flat #f)])
      (gc:first flat-addr))
    "not a pair")

  ;; dereferencing flat as cons
  (test/heap/exn (make-vector 1000)
    (let ([flat-addr (gc:alloc-flat #f)])
      (gc:rest flat-addr))
    "not a pair")

  ;; setting flat as cons
  (test/heap/exn (make-vector 1000)
    (let ([flat-addr (gc:alloc-flat #f)])
      (gc:set-first! flat-addr #t))
    "not a pair")

  ;; setting flat as cons
  (test/heap/exn (make-vector 1000)
    (let ([flat-addr (gc:alloc-flat #f)])
      (gc:set-rest! flat-addr #t))
    "not a pair")

  ;; getting code ptr from non closure
  (test/heap/exn (make-vector 1000)
    (let ([flat-addr (gc:alloc-flat #f)])
      (gc:closure-code-ptr flat-addr))
    "not a closure")

  ;; getting code ptr from non closure
  (test/heap/exn (make-vector 1000)
    (let ([flat-addr (gc:alloc-flat #f)])
      (gc:closure-env-ref flat-addr 1))
    "not a closure")

  ;; Successful dereference: flat
  (test/heap (make-vector 1000)
    (gc:deref (gc:alloc-flat #t))
    #t)

  ;; successful dereference: cons
  (test/heap (make-vector 1000)
    (gc:deref
     (gc:rest
      (gc:cons
       (simple-root (gc:alloc-flat 'first))
       (simple-root (gc:alloc-flat 'rest)))))
    'rest)

  (test/heap (make-vector 1000)
    (gc:deref
     (gc:first
      (gc:cons
       (simple-root (gc:alloc-flat 'first))
       (simple-root (gc:alloc-flat 'rest)))))
    'first)

  ;; successful alloc / deref closure
  (test/heap (make-vector 1000)
    (gc:closure-code-ptr
     (gc:closure 'dummy '()))
    'dummy)

  (test/heap (make-vector 1000)
    (gc:deref
     (gc:closure-env-ref
      (gc:closure
       'dummy
       (list (simple-root (gc:alloc-flat #f))))
      0))
    #f)

  ;; setting cons parts
  (test/heap (make-vector 1000)
    (let ([cons-loc
           (gc:cons
            (simple-root (gc:alloc-flat 'first))
            (simple-root (gc:alloc-flat 'rest)))])
      (gc:set-first! cons-loc (gc:alloc-flat 'mutated))
      (gc:deref (gc:first cons-loc)))
    'mutated)

  (test/heap (make-vector 1000)
    (let ([cons-loc
           (gc:cons
            (simple-root (gc:alloc-flat 'first))
            (simple-root (gc:alloc-flat 'rest)))])
      (gc:set-rest! cons-loc (gc:alloc-flat 'mutated))
      (gc:deref (gc:rest cons-loc)))
    'mutated)
  )

(module+ test
  ;; heap state after initial allocation
  (test/heap
      (make-vector 12 '?)
    (gc:alloc-flat #f)
    (current-heap)
    #(2 2 flat #f ? ?  ? ? ? ? ? ?))

  (test/heap
      (make-vector 18  '?)
    (gc:cons
     (simple-root (gc:alloc-flat #f))
     (simple-root (gc:alloc-flat #t)))
    (current-heap)
    #(2 7 flat #f flat #t cons 2 4 ? ? ? ? ? ? ? ? ?))

  (test/heap
      (make-vector 18  '?)
    (gc:closure
     'dummy
     (list (simple-root (gc:alloc-flat #f))))
    (current-heap)
    #(2 6 flat #f clos dummy 1 2 ? ? ? ? ? ? ? ? ? ?))
  )

(module+ test
  ;; heap state and roots after gc

  (test/heap
      (make-vector 12 '?)
    (define f1 (gc:alloc-flat #f))
    (with-roots (f1)
      (gc!)
      (cons (current-heap) (map read-root (get-root-set))))
    (cons
     #(7 2 forwarded 7 ? ? ? flat #f ? ? ?)
     '(7)))

  (test/heap
      (make-vector 18  '?)
    (define c1
      (gc:cons
       (simple-root (gc:alloc-flat #f))
       (simple-root (gc:alloc-flat #t))))
    (with-roots (c1)
      (gc!)
      (cons (current-heap) (map read-root (get-root-set))))
    (cons
     #(10 7 forwarded 13 forwarded 15 forwarded 10 4 ? cons 13 15 flat #f flat #t ?)
     '(10)))

  (test/heap
      (make-vector 18  '?)
    (define cl1
      (gc:closure 'dummy (list (simple-root (gc:alloc-flat #f)))))
    (with-roots (cl1)
      (gc!)
      (cons (current-heap) (map read-root (get-root-set))))
    (cons
     #(10 6 forwarded 14 forwarded 10 1 2 ? ? clos dummy 1 14 flat #f ? ?)
     '(10)))

  ;; Test for coverage of forwarded tags.
  (test/heap
      (make-vector 26 '?)
    (define c1
      (gc:cons
       (simple-root (gc:alloc-flat 2))
       (simple-root (gc:alloc-flat empty))))
    (define c2
      (gc:cons
       (simple-root (gc:alloc-flat 1))
       (simple-root c1)))
    ;; force both cons cells to be moved before starting update pass
    (with-roots (c1 c2)
      (gc!)
      (current-heap))
    #(14 12
         forwarded 20 forwarded 22 forwarded 14 4 forwarded 24 forwarded 17 6
         cons 20 22 cons 24 14 flat 2 flat () flat 1)
    )
  )

(module+ test
  (test/heap
      (make-vector 12 '?)
    (define f1 (gc:alloc-flat #f))
    (gc! (simple-root f1))
    (current-heap)
    #(7 2 forwarded 7 ? ? ? flat #f ? ? ?)))