#lang plai/gc2/collector
(print-only-errors #t)

;; 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))))]))

;(heap-size)>=(bitmap-words)+block-width*(bitmap-words)
(define (block-width) 8)
(define (bitmap-words)
  (quotient (heap-size)
            (add1 (block-width))))

(define (init-allocator)
  (for ([i (in-range (bitmap-words))])
    (heap-set! i 0)))

(module+ test
  (with-heap (make-vector 9 '?) (test (bitmap-words) 1))
  (with-heap (make-vector 8 '?) (test (bitmap-words) 0))
  (with-heap (make-vector 13 '?) (test (bitmap-words) 1))
  (with-heap (make-vector 21 '?) (test (bitmap-words) 2))
  (with-heap (make-vector 27 '?) (test (bitmap-words) 3)))

(module+ test
  (test/heap (make-vector 9 '?)
    (init-allocator)
    (current-heap)
    #(0 ? ? ? ? ? ? ? ?)))

(define (ones k) (sub1 (expt 2 k)))
(define (update-bits! loc how-many set?)
  (define (flip bits)
    (bitwise-xor bits (ones (block-width))))
  (let* ([addr (- loc (bitmap-words))]
         [block (quotient addr (block-width))]
         [index (- addr (* block (block-width)))]
         [diff
          (arithmetic-shift (ones how-many) index)]
         [current (heap-ref block)])
    (heap-set!
     block
     (if set? (bitwise-ior current diff)
         (bitwise-and current (flip diff))))))

;; use dynamic programming to find longest run of 0s
(define (ending-at bits pos acc best)
  (cond
    [(>= pos (block-width)) best]
    [(bitwise-bit-set? bits pos)
     (ending-at bits
                (add1 pos) 0
                (max acc best))]
    [else
     (ending-at bits
                (add1 pos) (add1 acc)
                (max (add1 acc) best))]))

;; memoize the gap finding
(define max-gap
  (let ([gap-table (make-vector (expt 2 (block-width))
                                #f)])
    (lambda (bits)
      (unless (vector-ref gap-table bits)
        (vector-set! gap-table bits
                     (ending-at bits 0 0 0)))
      (vector-ref gap-table bits))))

(module+ test
  (test (max-gap 0) (block-width))
  (test (max-gap 15) (- (block-width) 4))
  (test (max-gap (sub1 (expt 2 (block-width)))) 0)
  ;; specific to 8 bit width
  (test (max-gap 145) 3))

;; assumes there is a gap
;; assumes there is a gap
(define (first-fit block-bits new-bits)
  (define (loop mask offset)
    (cond
      [(>= offset (block-width)) (error 'first-fit "internal error")]
      [(zero? (bitwise-and block-bits mask)) offset]
      [else (loop (arithmetic-shift mask 1)  (add1 offset))]))
  (loop new-bits 0))

(define (find-free-space n)
  (define (loop i)
    (define bits (heap-ref i))
    (cond
      [(> n (block-width))
       (error 'find-free-space
              "allocation > ~a" (block-width))]
      [(>= i (bitmap-words)) #f]
      [(>= (max-gap bits) n)
       (+ (* (block-width) i)
          (first-fit bits (ones n)))]
      [else (loop (add1 i))]))
  (define offset (loop 0))
  (and offset (+ (bitmap-words) offset)))
(module+ test
  (test/heap
      (make-vector 10 '?)
    (find-free-space 4)
    (bitmap-words))

  (test/heap
      (make-vector 9 '?)
    (find-free-space 4)
    (bitmap-words))

  (test/heap
      (make-vector 9 '?)
    (find-free-space 8)
    (bitmap-words))

  (test/heap/exn
      (make-vector 9 '?)
    (find-free-space 9)
    "allocation > 8")
  )

(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)))
  (unless second
    (error 'alloc "out of memory"))
  (update-bits! second n #t) ;; CHANGED
  second)
(module+ test
  (test/heap
      (make-vector 9 '?)
    (malloc 2)
    (current-heap)
    #(3 ? ? ? ? ? ? ? ?))

  (test/heap/exn
      (make-vector 4 '?)
    (malloc 2)
    (malloc 2)
    "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))

(define (traverse/roots roots)
  (cond
    [(list? roots)
     (for-each traverse/roots roots)]
    [(root? roots)
     (traverse/loc (read-root roots))]
    [else
     (error 'traverse/roots
            "unexpected roots: ~a" roots)]))

(define (traverse/loc loc)
  (case (heap-ref loc)
    [(flat gray-flat cons gray-cons clos gray-clos) (void)]
    [(white-flat) (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)]))

(define-syntax-rule (for/bitmap (loc loop) expr ...)
  (for/bitmap/proc (lambda (loc loop) expr ...)))

(define (for/bitmap/proc action)
  (for ([block (in-range 0 (bitmap-words))]
        #:unless (zero? (heap-ref block)))
    (define start (+ (bitmap-words)
                           (* block (block-width))))
    (define (loop loc)
      (define index (- loc start))
      (cond
        [(>= index (block-width)) (void)]
        [(bitwise-bit-set?
          (heap-ref block) index)
         (action loc loop)]
        [else (loop (add1 loc))]))
    (loop start)))
; validate-heap : -> void?

(define (validate-heap)
  (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)))

  (for/bitmap (loc loop)
    (case (heap-ref loc)
      [(flat) (loop (+ loc 2))]
      [(cons)
       (valid-pointer? (heap-ref (+ loc 1)))
       (valid-pointer? (heap-ref (+ loc 2)))
       (loop (+ loc 3))]
      [(clos)
       (for ([j (in-range 0 (heap-ref (+ loc 2)))])
         (valid-pointer? (heap-ref (+ loc 3 j))))
       (loop (+ loc 3 (heap-ref (+ loc 2))))]
      [else (error 'validate-heap "unexpected at pos. ~a in ~a" loc
                   (current-heap))])))

(module+ test
  (with-heap
    (make-vector 9 '?)
    (init-allocator)
    (test (validate-heap) (void)))
  (with-heap
    (make-vector 9 '?)
    (init-allocator)
    (gc:alloc-flat 42)
    (test (validate-heap) (void)))
  (with-heap
    (make-vector 18 '?)
    (init-allocator)
    (gc:alloc-flat 42)
    (test (validate-heap) (void)))
  (with-heap
    (make-vector 18 '?)
    (init-allocator)
    (for ([i (in-range 7)])
      (gc:alloc-flat i))
    (test (validate-heap) (void)))
  )

(define (mark-white!)
  (for/bitmap (loc loop)
    (case (heap-ref loc)
      [(cons)
       (heap-set! loc 'white-cons)
       (loop (+ loc 3))]
      [(flat)
       (heap-set! loc 'white-flat)
       (loop (+ loc 2))]
      [(clos)
       (heap-set! loc 'white-clos)
       (loop (+ loc 3 (heap-ref (+ loc 2))))]
      [else (error 'mark-white!
                   "unexpected tag: ~a" loc)])))

(module+ test
  (with-heap
    (make-vector 9 '?)
    (init-allocator)
    (mark-white!)
    (test (current-heap)  #(0 ? ? ? ? ? ? ? ?)))

  (with-heap
    (make-vector 9 '?)
    (init-allocator)
    (gc:alloc-flat 42)
    (gc:alloc-flat 43)
    (mark-white!)
    (test (current-heap) #(15 white-flat 42 white-flat 43 ? ? ? ?)))

  (with-heap
    (make-vector 18 '?)
    (init-allocator)
    (for ([i (in-range 7)])
      (gc:alloc-flat i))
    (mark-white!)
    (test (current-heap)
          #(255 63 white-flat 0 white-flat 1 white-flat 2 white-flat 3 white-flat 4 white-flat 5 white-flat 6 ? ?)))
  )

(define (free-white!)
  (for/bitmap (loc loop)
    (define (free! width)
      (update-bits! loc width #f) (loop (+ loc width)))
    (case (heap-ref loc)
      [(white-clos) (free! (+ 3 (heap-ref (+ loc 2))))]
      [(clos) (loop (+ loc 3 (heap-ref (+ loc 2))))]
      [(white-flat) (free! 2)]
      [(flat) (loop (+ loc 2))]
      [(white-cons) (free! 3)]
      [(cons) (loop (+ loc 3))]
      [else (error 'free-white!
                   "bad tag at ~a" loc)])))

; 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)


(module+ test
  (test/heap
      (make-vector 9 '?)
    (gc:alloc-flat 1)
    (current-heap)
    #(3 flat 1 ? ? ? ? ? ?))

  (test/heap
      (make-vector 9 '?)
    (gc:alloc-flat 1)
    (find-free-space 2)
    3)

  (test/heap
      (make-vector 9 '?)
    (gc:alloc-flat 1)
    (gc:alloc-flat 2)
    (current-heap)
    #(15 flat 1 flat 2 ? ? ? ?))

  (with-heap (make-vector 9 '?)
    (init-allocator)
    (gc:cons (simple-root (gc:alloc-flat 'first))
             (simple-root (gc:alloc-flat 'rest)))
    (test (current-heap) #(127 flat first flat rest cons 1 3 ? )))

  (with-heap (make-vector 9 '?)
    (init-allocator)
    (gc:closure 'code (list (simple-root (gc:alloc-flat 'value))))
    (test (current-heap) #(63 flat value clos code 1 1 ? ?)))

  (with-heap (make-vector 9 '?)
    (init-allocator)
    (collect-garbage)
    (test (vector-copy (current-heap) 0 4) #(0 ? ? ?))
    (gc:alloc-flat 42)
    (test (vector-copy (current-heap) 0 6) #(3 flat 42 ? ? ?))
    (collect-garbage)
    (test (vector-copy (current-heap) 0 4) #(0 white-flat 42 ?))
    (gc:closure 'code (list (simple-root (gc:alloc-flat 'value))))
    (test (current-heap) #(63 flat value clos code 1 1 ? ?))
    (collect-garbage)
    (test (vector-copy (current-heap) 0 4) #(0 white-flat value white-clos))
    (gc:cons (simple-root (gc:alloc-flat 'first))
             (simple-root (gc:alloc-flat 'rest)))
    (test (current-heap)  #(127 flat first flat rest cons 1 3 ?))
    (collect-garbage)
    (test  (current-heap)  #(0 white-flat first white-flat rest white-cons 1 3 ?))
    ))

(module+ test

  (with-heap (make-vector 4)
    (init-allocator)
    (test/exn
     (gc:cons
      (simple-root (gc:alloc-flat #f))
      (simple-root (gc:alloc-flat #t)))
     "out of memory"))

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

  (with-heap (make-vector 1000)
    (init-allocator)
    (test/exn
     (let ([flat-addr (gc:alloc-flat #f)])
       (gc:first flat-addr))
     "not a cons"))

  (with-heap (make-vector 1000)
    (init-allocator)
    (test/exn
     (let ([flat-addr (gc:alloc-flat #f)])
       (gc:closure-code-ptr flat-addr))
     "not a closure"))

  (with-heap (make-vector 1000)
    (init-allocator)
    (test/exn
     (let ([flat-addr (gc:alloc-flat #f)])
       (gc:closure-env-ref flat-addr 0))
     "not a closure"))
  )