#lang plai/gc2/collector
;; free-list-node := free-2 next |  free-n next size
;; next := location? | #f
(define (fl:node? loc)
  (member (heap-ref loc) '(free-2 free-n)))

(define (fl:check loc)
  (unless (fl:node? loc)
    (error 'fl:check "not free list node ~a" loc)))

(define (fl:next loc)
  (fl:check loc)
  (heap-ref (+ 1 loc)))
(define (fl:set-next! prev loc)
  (when prev (fl:check prev))
  (heap-set! (if prev (+ prev 1) FREE-LIST) loc))

(define (fl:init! loc size next)
  (case size
    [(2) (heap-set! loc 'free-2)]
    [else
     (heap-set! loc 'free-n)
     (fl:set-length! loc size)])
  (fl:set-next! loc next))
(define (fl:set-length! loc len)
  (unless (eq? (heap-ref loc) 'free-n)
    (error 'fl:set-length! "illegal tag at ~a" loc))
  (heap-set! (+ 2 loc) len))
(define FREE-LIST 0) ; location of free list head
(define HEAP-START 1) ; where the user heap starts
(define (init-allocator)
  (heap-set! FREE-LIST HEAP-START) ; head of free list
  (fl:init! HEAP-START (- (heap-size) HEAP-START) #f))
(define (find-free-space size)
  (define (loop start prev)
    (define (split-current! loc free-size)
      (case free-size
        [(1) (delete-current!) (heap-set! loc 'free)]
        [else
         (fl:init! loc free-size (fl:next start))
         (fl:set-next! prev loc)]))
    (define (delete-current!)
      (fl:set-next! prev (fl:next start)))
     (case (heap-ref start)
       [(free-2)
        (cond
          [(= size 2) (delete-current!) start]
          [else (loop (fl:next start) start)])]
       [(free-n)
        (define length (heap-ref (+ start 2)))
        (cond
          [(= size length) (delete-current!) start]
          [(< size length)
           (split-current! (+ start size)
                           (- length size))
           start]
          [else (loop (fl:next start) start)])]
       [else (error 'find-free-space "wrong tag ~s at ~s"
                    (heap-ref start) start)])#;(@$\vdots$@))
(let ([head (heap-ref FREE-LIST)])
  (and head (loop head #f))))
; gc:flat? : location? -> boolean?
(define (gc:flat? addr)
  (equal? (heap-ref addr) 'flat))
; gc:deref location? -> heap-value?
(define (gc:deref addr)
  (unless (gc:flat? addr)
    (error 'gc:flat? "not a flat: ~a" addr))
  (heap-ref (+ addr 1)))

; gc:cons? : location? -> boolean?
(define (gc:cons? addr)
  (equal? (heap-ref addr) 'cons))
; gc:first : location? -> location?
(define (gc:first addr)
  (unless (gc:cons? addr)
    (error 'gc:first "not a cons: ~a" addr))
  (heap-ref (+ addr 1)))
; gc:rest : location? -> location?
(define (gc:rest addr)
  (unless (gc:cons? addr)
    (error 'gc:rest "not a cons: ~a" addr))
  (heap-ref (+ addr 2)))

; gc:set-first! : location? location? -> void?
(define (gc:set-first! addr v)
  (unless (gc:cons? addr)
    (error 'gc:set-first! "not a cons: ~a" addr))
  (heap-set! (+ addr 1) v))
; gc:set-rest! : location? location? -> void
(define (gc:set-rest! addr v)
  (unless (gc:cons? addr)
    (error 'gc:set-rest! "not a cons: ~a" addr))
  (heap-set! (+ addr 2) v))

; gc:closure? : location? -> boolean?
(define (gc:closure? addr)
  (equal? (heap-ref addr) 'clos))
; gc:closure-code-ptr : location? -> heap-value?
(define (gc:closure-code-ptr addr)
  (unless (gc:closure? addr)
    (error 'gc:closure-code-ptr "not a closure: ~a" addr))
  (heap-ref (+ addr 1)))
; gc:closure-env-ref : location? integer? -> location?
(define (gc:closure-env-ref addr i)
  (unless (gc:closure? addr)
    (error 'gc:closure-env-ref "not a closure: ~a" addr))
  (heap-ref (+ addr 3 i)))

; gc:alloc-flat : heap-value? -> location?
(define (gc:alloc-flat v)
  (define address (malloc 2))
  (heap-set! address 'flat)
  (heap-set! (+ 1 address) v)
  address)
; gc:cons : root? root? -> location?
(define (gc:cons v1 v2)
  (define address (malloc 3 v1 v2))
  (heap-set! address 'cons)
  (heap-set! (+ address 1) (read-root v1))
  (heap-set! (+ address 2) (read-root v2))
  address)

; gc:closure : heap-value? (vectorof location?) -> location?
(define (gc:closure code-ptr free-variables)
  (define address
    (malloc (+ 3 (length free-variables))
            free-variables #f))
  (heap-set! address 'clos)
  (heap-set! (+ address 1) code-ptr)
  (heap-set! (+ address 2) (length free-variables))
  (for ([i (in-range 0 (length free-variables))]
        [f (in-list free-variables)])
    (heap-set! (+ address 3 i) (read-root f)))
  address)


(define (malloc n . extra-roots)
  (define initial (find-free-space n))
  (unless initial
    (collect-garbage extra-roots))
  (define second (or initial (find-free-space n)))
  (if second second
      (error 'alloc "out of memory")))

; collect-garbage : roots? -> void?
(define (collect-garbage . extra-roots)
  (validate-heap)
  (mark-white!)
  (traverse/roots (get-root-set))
  (traverse/roots extra-roots)
  (free-white!)
  (validate-heap))

; validate-heap : -> void?
(define (validate-heap)
  (define (valid-pointer-or-#f p)
    (when p (valid-pointer? p)))
  (define (valid-pointer? p)
    (unless (< p (heap-size))
      (error 'validate-heap "pointer out of bounds ~a" p))
    (unless (member (heap-ref p) '(flat cons clos free-n free-2))
      (error 'validate-heap "pointer to non-tag ~a" p)))
  (let loop ([i HEAP-START])
    (when (< i (heap-size))
      (case (heap-ref i)
        [(flat) (loop (+ i 2))]
        [(cons)
         (valid-pointer? (heap-ref (+ i 1)))
         (valid-pointer? (heap-ref (+ i 2)))
         (loop (+ i 3))]
        [(clos)
         (for ([j (in-range 0 (heap-ref (+ i 2)))])
           (valid-pointer? (heap-ref (+ i 3 j))))
         (loop (+ i 3 (heap-ref (+ i 2))))]
        [(free-2)
         (valid-pointer-or-#f (heap-ref (+ i 1)))
         (loop (+ i 2))]
        [(free-n)
         (valid-pointer-or-#f (heap-ref (+ i 1)))
         (loop (+ i (heap-ref (+ i 2))))]
        [(free) (loop (+ i 1))]
        [else (error 'validate-heap "unexpected tag: ~a" i)]))))

; mark-white! : -> void?
(define (mark-white!)
  (let loop ([i HEAP-START])
    (when (< i (heap-size))
      (define tag (heap-ref i))
      (case tag
        [(cons)
         (heap-set! i 'white-cons)
         (loop (+ i 3))]
        [(flat)
         (heap-set! i 'white-flat)
         (loop (+ i 2))]
        [(clos)
         (heap-set! i 'white-clos)
         (loop (+ i 3 (heap-ref (+ i 2))))]
        [(free)
         (loop (+ i 1))]
        [(free-2)
         (loop (+ i 2))]
        [(free-n)
         (loop (+ i (heap-ref (+ i 2))))]
        [else (error 'mark-white!
                     "unexpected tag ~a at ~a" tag i)]))))


;; object-length : location -> number
(define (object-length loc)
  (define tag (heap-ref loc))
  (case tag
    [(free) 1]
    [(free-2) 2]
    [(free-n) (heap-ref (+ loc 2))]
    [(flat white-flat) 2]
    [(cons white-cons) 3]
    [(clos white-clos) (+ 3 (heap-ref (+ loc 2)))]
    [else (error 'object-length "wrong tag ~s at ~s" tag loc)]))
; traverse/roots : roots? -> void?
(define (traverse/roots roots)
  (cond [(list? roots)
         (for-each traverse/roots roots)]
        [(root? roots)
         (traverse/loc (read-root roots))]
        [(false? roots)
         (void)]
        [else
         (error 'traverse/roots
                "unexpected roots: ~a" roots)]))

(define (traverse/loc loc)
  (case (heap-ref loc)
    [(flat gray-flat) (void)]
    [(cons gray-cons) (void)]
    [(clos gray-clos) (void)]
    [(white-flat)
     ; can skip gray
     (heap-set! loc 'flat)]
    [(white-cons)
     (heap-set! loc 'gray-cons)
     (traverse/loc (heap-ref (+ loc 1)))
     (traverse/loc (heap-ref (+ loc 2)))
     (heap-set! loc 'cons)]
    [(white-clos)
     (heap-set! loc 'gray-clos)
     (for ([i (in-range 0 (heap-ref (+ loc 2)))])
       (traverse/loc (heap-ref (+ loc i 3))))
     (heap-set! loc 'clos)]
    [else (error 'traverse/loc "unexpected tag: ~a" loc)]))

(module+ test
  (print-only-errors #t)
  (with-heap (make-vector 6 #f)
    (init-allocator)
    (test (current-heap) #(1 free-n #f 5 #f #f))
    (gc:alloc-flat 42)
    (test (current-heap) #(3 flat 42 free-n #f 3))
    (gc:alloc-flat 43)
    (test (current-heap) #(#f flat 42 flat 43 free)))

  (with-heap (make-vector 10 #f)
    (init-allocator)
    (gc:cons (simple-root (gc:alloc-flat 'first))
             (simple-root (gc:alloc-flat 'rest)))
    (test (current-heap) #(8 flat first flat rest cons 1 3 free-2 #f)))

  (with-heap (make-vector 10 #f)
    (init-allocator)
    (gc:closure 'code (list (simple-root (gc:alloc-flat 'value))))
    (test (current-heap) #(7 flat value clos code 1 1 free-n #f 3)))

  (with-heap (make-vector 10 #f)
    (init-allocator)
    (collect-garbage #f #f)
    (test (vector-copy (current-heap) 0 4) #(1 free-n #f 9))
    (gc:alloc-flat 42)
    (test (vector-copy (current-heap) 0 6) #(3 flat 42 free-n #f 7))
    (collect-garbage #f #f)
    (test (vector-copy (current-heap) 0 4) #(1 free-n #f 9))
    (gc:closure 'code (list (simple-root (gc:alloc-flat 'value))))
    (test (vector-copy (current-heap) 0 10) #(7 flat value clos code 1 1 free-n #f 3))
    (collect-garbage #f #f)
    (test (vector-copy (current-heap) 0 4) #(1 free-n #f 9))
    (gc:cons (simple-root (gc:alloc-flat 'first))
             (simple-root (gc:alloc-flat 'rest)))
    (test (vector-copy (current-heap) 0 10) #(8 flat first flat rest cons 1 3 free-2 #f))
    (collect-garbage #f #f)
    (test (vector-copy (current-heap) 0 4) #(1 free-n #f 9))
    ))




;; Part 2
(module+ test
  (with-heap (make-vector 1000)
    (init-allocator)
    (let ([flat-addr (gc:alloc-flat #t)])
      (test (gc:flat? flat-addr) #t)
      (test (gc:cons? flat-addr) #f)
      (test (gc:deref flat-addr) #t)))
  )


;; Part 3 cons cells
;; first and rest
(module+ test
  (with-heap (make-vector 1000)
    (init-allocator)
    (let ([cons-loc
           (gc:cons
            (simple-root (gc:alloc-flat 'first))
            (simple-root (gc:alloc-flat 'rest)))])
      (test (gc:deref (gc:rest cons-loc)) 'rest)
      (test (gc:deref (gc:first cons-loc)) 'first)))

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

; part 4 closures
(module+ test
  (with-heap (make-vector 1000)
    (init-allocator)
    (let ([closure-loc
           (gc:closure 'code-pointer  (list (simple-root (gc:alloc-flat 'sekrit))))])
      (test (gc:deref (gc:closure-env-ref closure-loc 0)) 'sekrit)
      (test (gc:closure-code-ptr closure-loc) 'code-pointer)))
  )

(define (free-white!)
  (define (loop loc prev last-start spaces-so-far)
    (define (tag-of len)
      (case len [(1) 'free] [(2) 'free-2]
            [else 'free-n]))
    (define (write-free-record! where next)
      (heap-set! where (tag-of spaces-so-far))
      (when (>= spaces-so-far 2)
        (heap-set! (+ 1 where) next))
      (when (>= spaces-so-far 3)
        (heap-set! (+ 2 where) spaces-so-far))
        (fl:set-next! prev last-start))
    (define merging (and last-start spaces-so-far 
                         (> spaces-so-far 1)))
    (cond
      [(>= loc (heap-size))
       (when merging (write-free-record! last-start #f))]
      [else
       (define length (object-length loc))
       (case (heap-ref loc)
         [(flat cons clos)
          (when merging (write-free-record! last-start #f))
          (loop (+ loc length)
                (if merging last-start prev) #f #f)]
         [(white-flat white-cons white-clos
                      free free-2 free-n)
          (cond
            [(and last-start spaces-so-far)
             (loop (+ loc length) prev last-start
                   (+ spaces-so-far length))]
            [else (loop (+ loc length)
                        prev loc length)])]
         [else (error 'free-white! "wrong tag at ~a" loc)])#;(@$\vdots$@)])#;(@$\vdots$@))
  (loop HEAP-START #f #f #f))